diff --git a/README.md b/README.md index bbd4359..80c39f9 100644 --- a/README.md +++ b/README.md @@ -1,33 +1,78 @@ # ComfyUI Chunked Batch Nodes -Custom nodes for chunked batched image-to-image workflows. +Custom nodes for chunked batched image-to-image workflows and sequential temporal video cleanup. Included nodes: - **VAE Encode Batch Chunked** - **SamplerCustomAdvanced Chunked** - **VAE Decode Batch Chunked** +- **Flux Video Cleanup Temporal Advanced** ## Purpose -These nodes let you process large image/frame batches in smaller execution chunks while preserving batch order and `SamplerCustomAdvanced` behavior as closely as possible. +These nodes let you process large image/frame batches in smaller execution chunks while preserving the normal ComfyUI `SamplerCustomAdvanced` path as closely as possible. -Typical workflow: +Typical non-temporal workflow: `IMAGE batch -> VAE Encode Batch Chunked -> SamplerCustomAdvanced Chunked -> VAE Decode Batch Chunked` +Typical temporal workflow: + +`IMAGE batch -> Flux Video Cleanup Temporal Advanced` + +## Architecture + +The package now splits the reusable runtime into shared helpers: + +- `core_sampling.py` + - `encode_image_batch_chunked(...)` + - `decode_latent_batch_chunked(...)` + - `sample_latent_batch_chunked(...)` + - `sample_single_latent_temporal(...)` +- `nodes_batch_chunked.py` + - thin wrappers for the generic chunked batch nodes +- `nodes_temporal.py` + - the temporal recurrent video cleanup node + +This keeps `SamplerCustomAdvanced Chunked` as a generic memory-safety node instead of overloading it with frame-to-frame recurrence. + +## Flux Video Cleanup Temporal Advanced + +Supported temporal modes: + +- `off` + - pure batch path; no recurrence +- `prev_output_blend` + - blends the current input frame with the previous cleaned output before VAE encode +- `external_warped_prev` + - blends the current input frame with an externally warped previous image batch +- `internal_flow_warp` + - computes optical flow internally and warps the previous cleaned output before blending + +### Important behavior + +- In temporal modes, sampling runs one frame at a time so frame `t` can depend on frame `t-1`. +- `sample_chunk_size` is only used when `temporal_mode=off`. +- `reset_every_n` and `scene_cut_threshold` break recurrence explicitly. +- `lock_seed=True` supports fixed or sequential per-frame seeds based on the incoming `NOISE.seed`. +- `internal_flow_warp` uses OpenCV Farneback flow if `cv2` is available in the ComfyUI Python environment. + ## Notes -- `SamplerCustomAdvanced Chunked` reuses the normal `guider.sample(...)` path per chunk instead of reimplementing denoising logic. +- `SamplerCustomAdvanced Chunked` still reuses the normal `guider.sample(...)` path per chunk instead of reimplementing denoising logic. - When the input latent has no `batch_index` and chunking is required, the sampler node synthesizes a sequential internal `batch_index` so chunked random noise generation stays frame-stable across chunks. +- Temporal sequential sampling preserves frame-indexed noise variation when `lock_seed=False`, and intentionally pins batch index to `0` when `lock_seed=True`. - `noise_mask` batch slicing supports the common ComfyUI broadcast cases (`1` mask for the whole batch, or shorter masks that repeat across the full batch). - The sampler node supports OOM fallback by halving the chunk size down to `min_chunk_size`. +- VAE encode trims extra image channels beyond RGB before calling `vae.encode(...)`, which preserves compatibility with RGBA-style image batches without changing normal RGB inputs. ## Limitations -- This does **not** add temporal consistency. It only makes large frame batches tractable. -- Shared conditioning is the intended v1 use case. Per-frame conditioning batches inside the guider path are not explicitly rewritten by these nodes. -- The main target is regular latent batches used for Flux / img2img-style cleanup. Nested latent edge cases are handled conservatively but are not the primary tested path. +- The temporal node uses a pixel-space temporal prior before encode. It does not rewrite guider conditioning on a per-frame basis. +- `lock_seed=True` assumes the `NOISE` input exposes a mutable `seed` attribute, which matches ComfyUI's standard custom sampler noise objects. +- `internal_flow_warp` is a dependency-light built-in flow path, not a RAFT integration. +- Nested latent edge cases are still handled conservatively and are not the primary tested path. ## Installation diff --git a/__init__.py b/__init__.py index 5e560db..264630e 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,14 @@ -from .chunked_batch_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .nodes_batch_chunked import NODE_CLASS_MAPPINGS as BATCH_NODE_CLASS_MAPPINGS +from .nodes_batch_chunked import NODE_DISPLAY_NAME_MAPPINGS as BATCH_NODE_DISPLAY_NAME_MAPPINGS +from .nodes_temporal import NODE_CLASS_MAPPINGS as TEMPORAL_NODE_CLASS_MAPPINGS +from .nodes_temporal import NODE_DISPLAY_NAME_MAPPINGS as TEMPORAL_NODE_DISPLAY_NAME_MAPPINGS + +NODE_CLASS_MAPPINGS = {} +NODE_CLASS_MAPPINGS.update(BATCH_NODE_CLASS_MAPPINGS) +NODE_CLASS_MAPPINGS.update(TEMPORAL_NODE_CLASS_MAPPINGS) + +NODE_DISPLAY_NAME_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS.update(BATCH_NODE_DISPLAY_NAME_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(TEMPORAL_NODE_DISPLAY_NAME_MAPPINGS) __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/chunked_batch_nodes.py b/chunked_batch_nodes.py index f7f2e71..55a9ca6 100644 --- a/chunked_batch_nodes.py +++ b/chunked_batch_nodes.py @@ -1,331 +1,15 @@ -import gc -import math -from typing import Any, Dict, List, Tuple - -import torch - -import comfy.model_management -import comfy.nested_tensor -import comfy.sample -import comfy.utils -import latent_preview - - -LATENT_KEYS_WITH_BROADCAST_BATCH_ONE = {"noise_mask"} - - -def _soft_clear_cache() -> None: - gc.collect() - comfy.model_management.soft_empty_cache() - - -def _is_nested(value: Any) -> bool: - return getattr(value, "is_nested", False) - - -def _latent_batch_size(samples: Any) -> int: - if _is_nested(samples): - return len(samples.unbind()) - return int(samples.shape[0]) - - -def _slice_samples(samples: Any, start: int, end: int) -> Any: - if _is_nested(samples): - return comfy.nested_tensor.NestedTensor(samples.unbind()[start:end]) - return samples[start:end] - - -def _concat_samples(chunks: List[Any]) -> Any: - if not chunks: - raise ValueError("expected at least one chunk") - first = chunks[0] - if _is_nested(first): - tensors = [] - for chunk in chunks: - tensors.extend(chunk.unbind()) - return comfy.nested_tensor.NestedTensor(tensors) - return torch.cat(chunks, dim=0) - - -def _repeat_and_slice_batch_tensor(value: torch.Tensor, start: int, end: int, full_batch_size: int) -> torch.Tensor: - if value.shape[0] == full_batch_size: - return value[start:end] - if value.shape[0] == 1: - return value - if value.shape[0] < full_batch_size: - window_len = end - start - offset = start % value.shape[0] - reps = (math.ceil((offset + window_len) / value.shape[0]),) + ((1,) * (value.ndim - 1)) - expanded = value.repeat(reps) - return expanded[offset:offset + window_len] - return value[start:end] - - -def _slice_latent_value(key: str, value: Any, start: int, end: int, full_batch_size: int) -> Any: - if key == "samples": - return _slice_samples(value, start, end) - - if isinstance(value, torch.Tensor) and value.ndim > 0: - if value.shape[0] == full_batch_size: - return value[start:end] - if key in LATENT_KEYS_WITH_BROADCAST_BATCH_ONE and value.shape[0] <= full_batch_size: - return _repeat_and_slice_batch_tensor(value, start, end, full_batch_size) - return value - - if isinstance(value, list) and len(value) == full_batch_size: - return value[start:end] - - if isinstance(value, tuple) and len(value) == full_batch_size: - return value[start:end] - - return value - - -def _slice_latent_dict(latent: Dict[str, Any], start: int, end: int, full_batch_size: int) -> Dict[str, Any]: - return { - key: _slice_latent_value(key, value, start, end, full_batch_size) - for key, value in latent.items() - } - - -def _append_batch_index_if_needed(latent: Dict[str, Any], chunking_needed: bool) -> Tuple[Dict[str, Any], bool]: - if not chunking_needed or "batch_index" in latent: - return latent, False - - updated = latent.copy() - updated["batch_index"] = list(range(_latent_batch_size(latent["samples"]))) - return updated, True - - -def _process_denoised_chunk(guider: Any, x0_value: Any, sample_chunk: Any) -> Any: - x0_out = guider.model_patcher.model.process_latent_out(x0_value.cpu()) - if _is_nested(sample_chunk): - latent_shapes = [x.shape for x in sample_chunk.unbind()] - x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes)) - return x0_out - - -def _run_sampler_chunk(noise: Any, guider: Any, sampler: Any, sigmas: torch.Tensor, latent_chunk: Dict[str, Any]) -> Tuple[Any, Any, bool]: - latent_image = latent_chunk["samples"] - latent_image = comfy.sample.fix_empty_latent_channels( - guider.model_patcher, - latent_image, - latent_chunk.get("downscale_ratio_spacial", None), - ) - - working_latent = latent_chunk.copy() - working_latent["samples"] = latent_image - - noise_mask = working_latent.get("noise_mask", None) - x0_output: Dict[str, Any] = {} - callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output) - disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED - - samples = guider.sample( - noise.generate_noise(working_latent), - latent_image, - sampler, - sigmas, - denoise_mask=noise_mask, - callback=callback, - disable_pbar=disable_pbar, - seed=noise.seed, - ) - samples = samples.to(comfy.model_management.intermediate_device()) - - if "x0" in x0_output: - denoised = _process_denoised_chunk(guider, x0_output["x0"], samples) - return samples, denoised, True - - return samples, None, False - - -class VAEEncodeBatchChunked: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "pixels": ("IMAGE",), - "vae": ("VAE",), - "chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}), - "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), - } - } - - RETURN_TYPES = ("LATENT",) - RETURN_NAMES = ("latent",) - FUNCTION = "encode" - CATEGORY = "latent" - - def encode(self, pixels, vae, chunk_size, clear_cache_between_chunks=False): - batch_size = int(pixels.shape[0]) - if batch_size == 0: - raise ValueError("pixels batch must not be empty") - - encoded_chunks = [] - for start in range(0, batch_size, chunk_size): - end = min(start + chunk_size, batch_size) - encoded_chunks.append(vae.encode(pixels[start:end])) - if clear_cache_between_chunks and end < batch_size: - _soft_clear_cache() - - return ({"samples": _concat_samples(encoded_chunks)},) - - -class VAEDecodeBatchChunked: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "samples": ("LATENT",), - "vae": ("VAE",), - "chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}), - "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), - } - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("images",) - FUNCTION = "decode" - CATEGORY = "latent" - - def decode(self, samples, vae, chunk_size, clear_cache_between_chunks=False): - latent_samples = samples["samples"] - batch_size = _latent_batch_size(latent_samples) - if batch_size == 0: - raise ValueError("latent batch must not be empty") - - image_chunks = [] - for start in range(0, batch_size, chunk_size): - end = min(start + chunk_size, batch_size) - decoded = vae.decode(_slice_samples(latent_samples, start, end)) - if len(decoded.shape) == 5: - decoded = decoded.reshape(-1, decoded.shape[-3], decoded.shape[-2], decoded.shape[-1]) - image_chunks.append(decoded) - if clear_cache_between_chunks and end < batch_size: - _soft_clear_cache() - - return (torch.cat(image_chunks, dim=0),) - - -class SamplerCustomAdvancedChunked: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "noise": ("NOISE",), - "guider": ("GUIDER",), - "sampler": ("SAMPLER",), - "sigmas": ("SIGMAS",), - "latent_image": ("LATENT",), - "chunk_size": ("INT", {"default": 2, "min": 1, "max": 4096}), - "auto_reduce_on_oom": ("BOOLEAN", {"default": True}), - "min_chunk_size": ("INT", {"default": 1, "min": 1, "max": 4096}), - "clear_cache_on_retry": ("BOOLEAN", {"default": True, "advanced": True}), - "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), - } - } - - RETURN_TYPES = ("LATENT", "LATENT") - RETURN_NAMES = ("output", "denoised_output") - FUNCTION = "sample" - CATEGORY = "sampling/custom_sampling" - - def sample( - self, - noise, - guider, - sampler, - sigmas, - latent_image, - chunk_size, - auto_reduce_on_oom=True, - min_chunk_size=1, - clear_cache_on_retry=True, - clear_cache_between_chunks=False, - ): - original_latent = latent_image - full_samples = original_latent["samples"] - batch_size = _latent_batch_size(full_samples) - if batch_size == 0: - raise ValueError("latent batch must not be empty") - - effective_chunk_size = max(1, min(int(chunk_size), batch_size)) - min_chunk_size = max(1, min(int(min_chunk_size), effective_chunk_size)) - chunking_needed = batch_size > effective_chunk_size - - working_latent, synthetic_batch_index = _append_batch_index_if_needed(original_latent, chunking_needed) - - sample_chunks = [] - denoised_chunks = [] - x0_presence = [] - - start = 0 - while start < batch_size: - attempt_size = min(effective_chunk_size, batch_size - start) - end = start + attempt_size - latent_chunk = _slice_latent_dict(working_latent, start, end, batch_size) - retry_after_oom = False - - try: - samples_chunk, denoised_chunk, has_x0 = _run_sampler_chunk(noise, guider, sampler, sigmas, latent_chunk) - except Exception as exc: - if not auto_reduce_on_oom or not comfy.model_management.is_oom(exc) or effective_chunk_size <= min_chunk_size: - raise - - new_chunk_size = max(min_chunk_size, effective_chunk_size // 2) - if new_chunk_size == effective_chunk_size: - raise - - effective_chunk_size = new_chunk_size - if not synthetic_batch_index and "batch_index" not in original_latent and batch_size > effective_chunk_size: - working_latent, synthetic_batch_index = _append_batch_index_if_needed(original_latent, True) - retry_after_oom = True - - if retry_after_oom: - latent_chunk = None - if clear_cache_on_retry: - _soft_clear_cache() - continue - - sample_chunks.append(samples_chunk) - x0_presence.append(has_x0) - if has_x0: - denoised_chunks.append(denoised_chunk) - - start = end - if clear_cache_between_chunks and start < batch_size: - _soft_clear_cache() - - output_samples = _concat_samples(sample_chunks) - out = original_latent.copy() - out.pop("downscale_ratio_spacial", None) - out["samples"] = output_samples - - if synthetic_batch_index: - out.pop("batch_index", None) - - if any(x0_presence): - if not all(x0_presence): - raise RuntimeError("inconsistent x0 callback behavior across chunks") - out_denoised = original_latent.copy() - out_denoised["samples"] = _concat_samples(denoised_chunks) - if synthetic_batch_index: - out_denoised.pop("batch_index", None) - else: - out_denoised = out - - return (out, out_denoised) - - -NODE_CLASS_MAPPINGS = { - "VAEEncodeBatchChunked": VAEEncodeBatchChunked, - "VAEDecodeBatchChunked": VAEDecodeBatchChunked, - "SamplerCustomAdvancedChunked": SamplerCustomAdvancedChunked, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "VAEEncodeBatchChunked": "VAE Encode Batch Chunked", - "VAEDecodeBatchChunked": "VAE Decode Batch Chunked", - "SamplerCustomAdvancedChunked": "SamplerCustomAdvanced Chunked", -} +from .nodes_batch_chunked import ( + NODE_CLASS_MAPPINGS, + NODE_DISPLAY_NAME_MAPPINGS, + SamplerCustomAdvancedChunked, + VAEDecodeBatchChunked, + VAEEncodeBatchChunked, +) + +__all__ = [ + "NODE_CLASS_MAPPINGS", + "NODE_DISPLAY_NAME_MAPPINGS", + "SamplerCustomAdvancedChunked", + "VAEDecodeBatchChunked", + "VAEEncodeBatchChunked", +] diff --git a/core_sampling.py b/core_sampling.py new file mode 100644 index 0000000..cef0682 --- /dev/null +++ b/core_sampling.py @@ -0,0 +1,386 @@ +import contextlib +import gc +import math +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch + +import comfy.model_management +import comfy.nested_tensor +import comfy.sample +import comfy.utils +import latent_preview + + +LATENT_KEYS_WITH_BROADCAST_BATCH_ONE = {"noise_mask"} + + +@dataclass(frozen=True) +class SampleChunkConfig: + chunk_size: int + auto_reduce_on_oom: bool = True + min_chunk_size: int = 1 + clear_cache_on_retry: bool = True + clear_cache_between_chunks: bool = False + + +def soft_clear_cache() -> None: + gc.collect() + comfy.model_management.soft_empty_cache() + + +def _is_nested(value: Any) -> bool: + return getattr(value, "is_nested", False) + + +def _trim_vae_input_channels(pixels: torch.Tensor) -> torch.Tensor: + if pixels.ndim >= 4 and pixels.shape[-1] > 3: + return pixels[..., :3] + return pixels + + +def _latent_batch_size(samples: Any) -> int: + if _is_nested(samples): + return len(samples.unbind()) + return int(samples.shape[0]) + + +def _slice_samples(samples: Any, start: int, end: int) -> Any: + if _is_nested(samples): + return comfy.nested_tensor.NestedTensor(samples.unbind()[start:end]) + return samples[start:end] + + +def _concat_samples(chunks: List[Any]) -> Any: + if not chunks: + raise ValueError("expected at least one chunk") + first = chunks[0] + if _is_nested(first): + tensors = [] + for chunk in chunks: + tensors.extend(chunk.unbind()) + return comfy.nested_tensor.NestedTensor(tensors) + return torch.cat(chunks, dim=0) + + +def _repeat_and_slice_batch_tensor( + value: torch.Tensor, + start: int, + end: int, + full_batch_size: int, +) -> torch.Tensor: + if value.shape[0] == full_batch_size: + return value[start:end] + if value.shape[0] == 1: + return value + if value.shape[0] < full_batch_size: + window_len = end - start + offset = start % value.shape[0] + reps = (math.ceil((offset + window_len) / value.shape[0]),) + ((1,) * (value.ndim - 1)) + expanded = value.repeat(reps) + return expanded[offset : offset + window_len] + return value[start:end] + + +def _slice_latent_value(key: str, value: Any, start: int, end: int, full_batch_size: int) -> Any: + if key == "samples": + return _slice_samples(value, start, end) + + if isinstance(value, torch.Tensor) and value.ndim > 0: + if value.shape[0] == full_batch_size: + return value[start:end] + if key in LATENT_KEYS_WITH_BROADCAST_BATCH_ONE and value.shape[0] <= full_batch_size: + return _repeat_and_slice_batch_tensor(value, start, end, full_batch_size) + return value + + if isinstance(value, list) and len(value) == full_batch_size: + return value[start:end] + + if isinstance(value, tuple) and len(value) == full_batch_size: + return value[start:end] + + return value + + +def _slice_latent_dict(latent: Dict[str, Any], start: int, end: int, full_batch_size: int) -> Dict[str, Any]: + return { + key: _slice_latent_value(key, value, start, end, full_batch_size) + for key, value in latent.items() + } + + +def _append_batch_index_if_needed( + latent: Dict[str, Any], + chunking_needed: bool, + forced_batch_index: Optional[Sequence[int]] = None, +) -> Tuple[Dict[str, Any], bool]: + if forced_batch_index is not None: + if len(forced_batch_index) != _latent_batch_size(latent["samples"]): + raise ValueError("forced batch_index length must match latent batch size") + updated = latent.copy() + updated["batch_index"] = list(forced_batch_index) + existing = latent.get("batch_index") + return updated, existing != updated["batch_index"] + + if not chunking_needed or "batch_index" in latent: + return latent, False + + updated = latent.copy() + updated["batch_index"] = list(range(_latent_batch_size(latent["samples"]))) + return updated, True + + +@contextlib.contextmanager +def _temporary_noise_seed(noise: Any, seed_override: Optional[int]): + if seed_override is None: + yield noise + return + + if not hasattr(noise, "seed"): + raise TypeError("noise seed override requires a NOISE input with a mutable 'seed' attribute") + + original_seed = noise.seed + try: + noise.seed = seed_override + except Exception as exc: + raise TypeError("noise seed override requires a mutable 'seed' attribute on the NOISE input") from exc + try: + yield noise + finally: + noise.seed = original_seed + + +def _process_denoised_chunk(guider: Any, x0_value: Any, sample_chunk: Any) -> Any: + x0_out = guider.model_patcher.model.process_latent_out(x0_value.cpu()) + if _is_nested(sample_chunk): + latent_shapes = [x.shape for x in sample_chunk.unbind()] + x0_out = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0_out, latent_shapes)) + return x0_out + + +def _run_sampler_chunk( + noise: Any, + guider: Any, + sampler: Any, + sigmas: torch.Tensor, + latent_chunk: Dict[str, Any], + seed_override: Optional[int] = None, +) -> Tuple[Any, Any, bool]: + latent_image = latent_chunk["samples"] + latent_image = comfy.sample.fix_empty_latent_channels( + guider.model_patcher, + latent_image, + latent_chunk.get("downscale_ratio_spacial", None), + ) + + working_latent = latent_chunk.copy() + working_latent["samples"] = latent_image + + noise_mask = working_latent.get("noise_mask", None) + x0_output: Dict[str, Any] = {} + callback = latent_preview.prepare_callback(guider.model_patcher, sigmas.shape[-1] - 1, x0_output) + disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED + + with _temporary_noise_seed(noise, seed_override) as seeded_noise: + samples = guider.sample( + seeded_noise.generate_noise(working_latent), + latent_image, + sampler, + sigmas, + denoise_mask=noise_mask, + callback=callback, + disable_pbar=disable_pbar, + seed=seeded_noise.seed, + ) + samples = samples.to(comfy.model_management.intermediate_device()) + + if "x0" in x0_output: + denoised = _process_denoised_chunk(guider, x0_output["x0"], samples) + return samples, denoised, True + + return samples, None, False + + +def encode_image_batch_chunked( + pixels: torch.Tensor, + vae: Any, + chunk_size: int, + clear_cache_between_chunks: bool = False, +) -> Dict[str, Any]: + batch_size = int(pixels.shape[0]) + if batch_size == 0: + raise ValueError("pixels batch must not be empty") + + encoded_chunks = [] + step = int(chunk_size) + for start in range(0, batch_size, step): + end = min(start + step, batch_size) + encoded_chunks.append(vae.encode(_trim_vae_input_channels(pixels[start:end]))) + if clear_cache_between_chunks and end < batch_size: + soft_clear_cache() + + return {"samples": _concat_samples(encoded_chunks)} + + +def decode_latent_batch_chunked( + samples: Dict[str, Any], + vae: Any, + chunk_size: int, + clear_cache_between_chunks: bool = False, +) -> torch.Tensor: + latent_samples = samples["samples"] + batch_size = _latent_batch_size(latent_samples) + if batch_size == 0: + raise ValueError("latent batch must not be empty") + + image_chunks = [] + step = int(chunk_size) + for start in range(0, batch_size, step): + end = min(start + step, batch_size) + decoded = vae.decode(_slice_samples(latent_samples, start, end)) + if len(decoded.shape) == 5: + decoded = decoded.reshape(-1, decoded.shape[-3], decoded.shape[-2], decoded.shape[-1]) + image_chunks.append(decoded) + if clear_cache_between_chunks and end < batch_size: + soft_clear_cache() + + return torch.cat(image_chunks, dim=0) + + +def sample_latent_batch_chunked( + noise: Any, + guider: Any, + sampler: Any, + sigmas: torch.Tensor, + latent_image: Dict[str, Any], + config: SampleChunkConfig, + noise_seed_override: Optional[int] = None, + forced_batch_index: Optional[Sequence[int]] = None, +) -> Tuple[Dict[str, Any], Dict[str, Any]]: + original_latent = latent_image + full_samples = original_latent["samples"] + batch_size = _latent_batch_size(full_samples) + if batch_size == 0: + raise ValueError("latent batch must not be empty") + + effective_chunk_size = max(1, min(int(config.chunk_size), batch_size)) + min_chunk_size = max(1, min(int(config.min_chunk_size), effective_chunk_size)) + chunking_needed = batch_size > effective_chunk_size + + working_latent, synthetic_batch_index = _append_batch_index_if_needed( + original_latent, + chunking_needed, + forced_batch_index=forced_batch_index, + ) + + sample_chunks = [] + denoised_chunks = [] + x0_presence = [] + + start = 0 + while start < batch_size: + attempt_size = min(effective_chunk_size, batch_size - start) + end = start + attempt_size + latent_chunk = _slice_latent_dict(working_latent, start, end, batch_size) + retry_after_oom = False + + try: + samples_chunk, denoised_chunk, has_x0 = _run_sampler_chunk( + noise, + guider, + sampler, + sigmas, + latent_chunk, + seed_override=noise_seed_override, + ) + except Exception as exc: + if not config.auto_reduce_on_oom or not comfy.model_management.is_oom(exc) or effective_chunk_size <= min_chunk_size: + raise + + new_chunk_size = max(min_chunk_size, effective_chunk_size // 2) + if new_chunk_size == effective_chunk_size: + raise + + effective_chunk_size = new_chunk_size + if ( + forced_batch_index is None + and not synthetic_batch_index + and "batch_index" not in original_latent + and batch_size > effective_chunk_size + ): + working_latent, synthetic_batch_index = _append_batch_index_if_needed(original_latent, True) + retry_after_oom = True + + if retry_after_oom: + if config.clear_cache_on_retry: + soft_clear_cache() + continue + + sample_chunks.append(samples_chunk) + x0_presence.append(has_x0) + if has_x0: + denoised_chunks.append(denoised_chunk) + + start = end + if config.clear_cache_between_chunks and start < batch_size: + soft_clear_cache() + + output_samples = _concat_samples(sample_chunks) + out = original_latent.copy() + out.pop("downscale_ratio_spacial", None) + out["samples"] = output_samples + + if synthetic_batch_index and forced_batch_index is None: + out.pop("batch_index", None) + + if any(x0_presence): + if not all(x0_presence): + raise RuntimeError("inconsistent x0 callback behavior across chunks") + out_denoised = original_latent.copy() + out_denoised["samples"] = _concat_samples(denoised_chunks) + if synthetic_batch_index and forced_batch_index is None: + out_denoised.pop("batch_index", None) + else: + out_denoised = out + + return out, out_denoised + + +def sample_single_latent_temporal( + noise: Any, + guider: Any, + sampler: Any, + sigmas: torch.Tensor, + latent_image: Dict[str, Any], + noise_seed_override: Optional[int] = None, + forced_batch_index_value: int = 0, + clear_cache_between_chunks: bool = False, +) -> Tuple[Dict[str, Any], Dict[str, Any]]: + return sample_latent_batch_chunked( + noise, + guider, + sampler, + sigmas, + latent_image, + config=SampleChunkConfig( + chunk_size=1, + auto_reduce_on_oom=False, + min_chunk_size=1, + clear_cache_on_retry=True, + clear_cache_between_chunks=clear_cache_between_chunks, + ), + noise_seed_override=noise_seed_override, + forced_batch_index=[int(forced_batch_index_value)], + ) + + +__all__ = [ + "LATENT_KEYS_WITH_BROADCAST_BATCH_ONE", + "SampleChunkConfig", + "decode_latent_batch_chunked", + "encode_image_batch_chunked", + "sample_latent_batch_chunked", + "sample_single_latent_temporal", + "soft_clear_cache", +] diff --git a/nodes_batch_chunked.py b/nodes_batch_chunked.py new file mode 100644 index 0000000..fa54aed --- /dev/null +++ b/nodes_batch_chunked.py @@ -0,0 +1,127 @@ +from .core_sampling import ( + SampleChunkConfig, + decode_latent_batch_chunked, + encode_image_batch_chunked, + sample_latent_batch_chunked, +) + + +class VAEEncodeBatchChunked: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "pixels": ("IMAGE",), + "vae": ("VAE",), + "chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}), + "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), + } + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "encode" + CATEGORY = "latent" + + def encode(self, pixels, vae, chunk_size, clear_cache_between_chunks=False): + return ( + encode_image_batch_chunked( + pixels, + vae, + chunk_size, + clear_cache_between_chunks=clear_cache_between_chunks, + ), + ) + + +class VAEDecodeBatchChunked: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "samples": ("LATENT",), + "vae": ("VAE",), + "chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}), + "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "decode" + CATEGORY = "latent" + + def decode(self, samples, vae, chunk_size, clear_cache_between_chunks=False): + return ( + decode_latent_batch_chunked( + samples, + vae, + chunk_size, + clear_cache_between_chunks=clear_cache_between_chunks, + ), + ) + + +class SamplerCustomAdvancedChunked: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "noise": ("NOISE",), + "guider": ("GUIDER",), + "sampler": ("SAMPLER",), + "sigmas": ("SIGMAS",), + "latent_image": ("LATENT",), + "chunk_size": ("INT", {"default": 2, "min": 1, "max": 4096}), + "auto_reduce_on_oom": ("BOOLEAN", {"default": True}), + "min_chunk_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "clear_cache_on_retry": ("BOOLEAN", {"default": True, "advanced": True}), + "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), + } + } + + RETURN_TYPES = ("LATENT", "LATENT") + RETURN_NAMES = ("output", "denoised_output") + FUNCTION = "sample" + CATEGORY = "sampling/custom_sampling" + + def sample( + self, + noise, + guider, + sampler, + sigmas, + latent_image, + chunk_size, + auto_reduce_on_oom=True, + min_chunk_size=1, + clear_cache_on_retry=True, + clear_cache_between_chunks=False, + ): + return sample_latent_batch_chunked( + noise, + guider, + sampler, + sigmas, + latent_image, + config=SampleChunkConfig( + chunk_size=chunk_size, + auto_reduce_on_oom=auto_reduce_on_oom, + min_chunk_size=min_chunk_size, + clear_cache_on_retry=clear_cache_on_retry, + clear_cache_between_chunks=clear_cache_between_chunks, + ), + ) + + +NODE_CLASS_MAPPINGS = { + "VAEEncodeBatchChunked": VAEEncodeBatchChunked, + "VAEDecodeBatchChunked": VAEDecodeBatchChunked, + "SamplerCustomAdvancedChunked": SamplerCustomAdvancedChunked, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "VAEEncodeBatchChunked": "VAE Encode Batch Chunked", + "VAEDecodeBatchChunked": "VAE Decode Batch Chunked", + "SamplerCustomAdvancedChunked": "SamplerCustomAdvanced Chunked", +} diff --git a/nodes_temporal.py b/nodes_temporal.py new file mode 100644 index 0000000..b5f4b53 --- /dev/null +++ b/nodes_temporal.py @@ -0,0 +1,518 @@ +from typing import Any, Optional, Tuple + +import torch +import torch.nn.functional as F + +from .core_sampling import ( + SampleChunkConfig, + decode_latent_batch_chunked, + encode_image_batch_chunked, + sample_latent_batch_chunked, + sample_single_latent_temporal, + soft_clear_cache, +) + +try: + import cv2 # type: ignore +except Exception: # pragma: no cover + cv2 = None + + +def _frame_batch_size(images: torch.Tensor) -> int: + if images.ndim != 4: + raise ValueError(f"expected IMAGE tensor with shape [B,H,W,C], got {tuple(images.shape)}") + return int(images.shape[0]) + + +def _single_frame(images: torch.Tensor, index: int) -> torch.Tensor: + return images[index : index + 1] + + +def _validate_optional_frame_batch(name: str, value: Optional[torch.Tensor], batch_size: int) -> None: + if value is None: + return + actual = _frame_batch_size(value) + if actual not in {1, batch_size, max(1, batch_size - 1)}: + raise ValueError( + f"{name} batch size must be 1, the full frame count, or frame_count - 1; " + f"got {actual} for {batch_size} frames" + ) + + +def _extract_optional_frame( + name: str, + value: Optional[torch.Tensor], + frame_index: int, + batch_size: int, +) -> Optional[torch.Tensor]: + if value is None: + return None + + actual = _frame_batch_size(value) + if actual == 1: + return value + if actual == batch_size: + return _single_frame(value, frame_index) + if actual == batch_size - 1: + if frame_index == 0: + return None + return _single_frame(value, frame_index - 1) + raise ValueError(f"{name} batch size mismatch") + + +def _mask_batch_size(mask: torch.Tensor) -> int: + if mask.ndim == 2: + return 1 + if mask.ndim == 3: + return int(mask.shape[0]) + if mask.ndim == 4 and mask.shape[-1] == 1: + return int(mask.shape[0]) + raise ValueError(f"unsupported MASK tensor shape {tuple(mask.shape)}") + + +def _validate_optional_mask_batch(name: str, value: Optional[torch.Tensor], batch_size: int) -> None: + if value is None: + return + actual = _mask_batch_size(value) + if actual not in {1, batch_size, max(1, batch_size - 1)}: + raise ValueError( + f"{name} batch size must be 1, the full frame count, or frame_count - 1; " + f"got {actual} for {batch_size} frames" + ) + + +def _extract_optional_mask( + name: str, + value: Optional[torch.Tensor], + frame_index: int, + batch_size: int, +) -> Optional[torch.Tensor]: + if value is None: + return None + + if value.ndim == 2: + return value.unsqueeze(0) + + if value.ndim == 3: + actual = int(value.shape[0]) + if actual == 1: + return value + if actual == batch_size: + return value[frame_index : frame_index + 1] + if actual == batch_size - 1: + if frame_index == 0: + return None + return value[frame_index - 1 : frame_index] + + if value.ndim == 4 and value.shape[-1] == 1: + actual = int(value.shape[0]) + if actual == 1: + return value[..., 0] + if actual == batch_size: + return value[frame_index : frame_index + 1, ..., 0] + if actual == batch_size - 1: + if frame_index == 0: + return None + return value[frame_index - 1 : frame_index, ..., 0] + + raise ValueError(f"{name} batch size mismatch") + + +def _mask_to_image_weight(mask: Optional[torch.Tensor], reference_image: torch.Tensor) -> Optional[torch.Tensor]: + if mask is None: + return None + + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 4 and mask.shape[-1] == 1: + mask = mask[..., 0] + if mask.ndim != 3: + raise ValueError(f"expected mask with shape [B,H,W], got {tuple(mask.shape)}") + + mask = mask.to(device=reference_image.device, dtype=reference_image.dtype) + if mask.shape[0] == 1 and reference_image.shape[0] != 1: + mask = mask.expand(reference_image.shape[0], -1, -1) + + return mask.unsqueeze(-1).clamp(0.0, 1.0) + + +def _ensure_same_image_shape(current: torch.Tensor, previous: torch.Tensor, context: str) -> torch.Tensor: + if previous.shape == current.shape: + return previous.to(device=current.device, dtype=current.dtype) + + if previous.shape[0] != current.shape[0] or previous.shape[-1] != current.shape[-1]: + raise ValueError( + f"{context} has incompatible IMAGE shape {tuple(previous.shape)} " + f"for current frame {tuple(current.shape)}" + ) + + resized = previous.permute(0, 3, 1, 2) + resized = F.interpolate(resized, size=current.shape[1:3], mode="bilinear", align_corners=False) + return resized.permute(0, 2, 3, 1).to(device=current.device, dtype=current.dtype) + + +def _frame_difference_score(previous_frame: torch.Tensor, current_frame: torch.Tensor) -> float: + previous = _ensure_same_image_shape(current_frame, previous_frame, "previous input") + return float(torch.mean(torch.abs(current_frame - previous)).item()) + + +def _compute_reset( + frame_index: int, + previous_input: Optional[torch.Tensor], + current_input: torch.Tensor, + reset_every_n: int, + scene_cut_threshold: float, +) -> bool: + if frame_index == 0 or previous_input is None: + return True + if reset_every_n > 0 and frame_index % reset_every_n == 0: + return True + if scene_cut_threshold > 0.0 and _frame_difference_score(previous_input, current_input) >= scene_cut_threshold: + return True + return False + + +def _blend_images( + current: torch.Tensor, + prior: torch.Tensor, + strength: float, + confidence: Optional[torch.Tensor] = None, +) -> torch.Tensor: + prior = _ensure_same_image_shape(current, prior, "temporal prior") + alpha = torch.full_like(current[..., :1], float(strength)) + if confidence is not None: + alpha = alpha * _mask_to_image_weight(confidence, current) + alpha = alpha.clamp(0.0, 1.0) + return current * (1.0 - alpha) + prior * alpha + + +def _cv2_require() -> Any: + if cv2 is None: + raise RuntimeError("internal_flow_warp requires OpenCV (cv2) in the ComfyUI Python environment") + return cv2 + + +def _image_to_numpy_uint8(frame: torch.Tensor): + return frame.detach().cpu().clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8).numpy() + + +def _compute_backward_flow(current_frame: torch.Tensor, previous_frame: torch.Tensor): + cv2_mod = _cv2_require() + current_np = _image_to_numpy_uint8(current_frame[0]) + previous_np = _image_to_numpy_uint8(previous_frame[0]) + current_gray = cv2_mod.cvtColor(current_np, cv2_mod.COLOR_RGB2GRAY) + previous_gray = cv2_mod.cvtColor(previous_np, cv2_mod.COLOR_RGB2GRAY) + return cv2_mod.calcOpticalFlowFarneback( + current_gray, + previous_gray, + None, + pyr_scale=0.5, + levels=3, + winsize=15, + iterations=3, + poly_n=5, + poly_sigma=1.2, + flags=0, + ) + + +def _warp_with_flow(image: torch.Tensor, backward_flow) -> torch.Tensor: + cv2_mod = _cv2_require() + import numpy as np + + image_np = image[0].detach().cpu().numpy().astype("float32", copy=False) + height, width = backward_flow.shape[:2] + grid_x, grid_y = np.meshgrid(range(width), range(height)) + map_x = (grid_x.astype("float32") + backward_flow[..., 0]).astype("float32") + map_y = (grid_y.astype("float32") + backward_flow[..., 1]).astype("float32") + warped = cv2_mod.remap( + image_np, + map_x, + map_y, + interpolation=cv2_mod.INTER_LINEAR, + borderMode=cv2_mod.BORDER_REPLICATE, + ) + return torch.from_numpy(warped).unsqueeze(0).to(device=image.device, dtype=image.dtype) + + +def _photometric_flow_confidence( + current_frame: torch.Tensor, + previous_input_frame: torch.Tensor, + backward_flow, +) -> torch.Tensor: + warped_previous_input = _warp_with_flow(previous_input_frame, backward_flow) + error = torch.mean(torch.abs(current_frame - warped_previous_input), dim=-1, keepdim=True) + return (1.0 - error).clamp(0.0, 1.0) + + +def _resolve_temporal_prior( + temporal_mode: str, + frame_index: int, + current_frame: torch.Tensor, + previous_input: Optional[torch.Tensor], + previous_output: Optional[torch.Tensor], + warped_previous_images: Optional[torch.Tensor], + flow_confidence_mask: Optional[torch.Tensor], + flow_confidence_scale: float, + temporal_strength: float, + batch_size: int, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if temporal_mode == "off" or frame_index == 0 or previous_output is None: + return current_frame, None + + external_mask = _extract_optional_mask("flow_confidence", flow_confidence_mask, frame_index, batch_size) + if external_mask is not None: + external_mask = external_mask * float(flow_confidence_scale) + elif flow_confidence_scale < 1.0: + external_mask = torch.full( + (1, current_frame.shape[1], current_frame.shape[2]), + float(flow_confidence_scale), + device=current_frame.device, + dtype=current_frame.dtype, + ) + + if temporal_mode == "prev_output_blend": + return _blend_images(current_frame, previous_output, temporal_strength), None + + if temporal_mode == "external_warped_prev": + warped_frame = _extract_optional_frame( + "warped_previous_images", + warped_previous_images, + frame_index, + batch_size, + ) + if warped_frame is None: + raise ValueError("external_warped_prev requires warped_previous_images after the first frame") + return _blend_images(current_frame, warped_frame, temporal_strength, confidence=external_mask), warped_frame + + if temporal_mode == "internal_flow_warp": + if previous_input is None: + return current_frame, None + previous_input = _ensure_same_image_shape(current_frame, previous_input, "previous input") + previous_output = _ensure_same_image_shape(current_frame, previous_output, "previous output") + backward_flow = _compute_backward_flow(current_frame, previous_input) + warped_previous_output = _warp_with_flow(previous_output, backward_flow) + confidence = _photometric_flow_confidence(current_frame, previous_input, backward_flow) + if external_mask is not None: + confidence = confidence * _mask_to_image_weight(external_mask, current_frame) + return ( + _blend_images(current_frame, warped_previous_output, temporal_strength, confidence=confidence), + warped_previous_output, + ) + + raise ValueError(f"unsupported temporal_mode: {temporal_mode}") + + +def _seed_for_frame(noise: Any, frame_index: int, lock_seed: bool, seed_stride_mode: str) -> Optional[int]: + if not lock_seed: + return None + if not hasattr(noise, "seed"): + raise TypeError("lock_seed=True requires a NOISE input with a 'seed' attribute") + + base_seed = int(noise.seed) + if seed_stride_mode == "fixed": + return base_seed + if seed_stride_mode == "sequential": + return base_seed + frame_index + raise ValueError(f"unsupported seed_stride_mode: {seed_stride_mode}") + + +def _batch_index_for_frame(frame_index: int, lock_seed: bool) -> int: + if lock_seed: + return 0 + return int(frame_index) + + +class FluxVideoCleanupTemporalAdvanced: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "vae": ("VAE",), + "noise": ("NOISE",), + "guider": ("GUIDER",), + "sampler": ("SAMPLER",), + "sigmas": ("SIGMAS",), + "temporal_mode": (("off", "prev_output_blend", "external_warped_prev", "internal_flow_warp"),), + "temporal_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.01}), + "reset_every_n": ("INT", {"default": 0, "min": 0, "max": 4096}), + "scene_cut_threshold": ("FLOAT", {"default": 0.18, "min": 0.0, "max": 1.0, "step": 0.01}), + "flow_confidence_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "advanced": True}), + "lock_seed": ("BOOLEAN", {"default": True}), + "seed_stride_mode": (("fixed", "sequential"),), + "encode_chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}), + "sample_chunk_size": ( + "INT", + { + "default": 1, + "min": 1, + "max": 4096, + "tooltip": "Only used when temporal_mode is off. Temporal recurrence forces per-frame sampling.", + }, + ), + "decode_chunk_size": ("INT", {"default": 8, "min": 1, "max": 4096}), + "auto_reduce_on_oom": ("BOOLEAN", {"default": True}), + "min_chunk_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "clear_cache_on_retry": ("BOOLEAN", {"default": True, "advanced": True}), + "clear_cache_between_chunks": ("BOOLEAN", {"default": False, "advanced": True}), + }, + "optional": { + "warped_previous_images": ("IMAGE",), + "flow_confidence": ("MASK",), + }, + } + + RETURN_TYPES = ("IMAGE", "LATENT", "LATENT", "MASK") + RETURN_NAMES = ("images", "latents", "denoised_latents", "reset_mask") + FUNCTION = "cleanup" + CATEGORY = "sampling/custom_sampling" + + def cleanup( + self, + images, + vae, + noise, + guider, + sampler, + sigmas, + temporal_mode, + temporal_strength, + reset_every_n, + scene_cut_threshold, + flow_confidence_scale, + lock_seed, + seed_stride_mode, + encode_chunk_size, + sample_chunk_size, + decode_chunk_size, + auto_reduce_on_oom=True, + min_chunk_size=1, + clear_cache_on_retry=True, + clear_cache_between_chunks=False, + warped_previous_images=None, + flow_confidence=None, + ): + batch_size = _frame_batch_size(images) + if batch_size == 0: + raise ValueError("images batch must not be empty") + + _validate_optional_frame_batch("warped_previous_images", warped_previous_images, batch_size) + _validate_optional_mask_batch("flow_confidence", flow_confidence, batch_size) + + if temporal_mode == "off": + latents = encode_image_batch_chunked( + images, + vae, + encode_chunk_size, + clear_cache_between_chunks=clear_cache_between_chunks, + ) + sampled, denoised = sample_latent_batch_chunked( + noise, + guider, + sampler, + sigmas, + latents, + config=SampleChunkConfig( + chunk_size=sample_chunk_size, + auto_reduce_on_oom=auto_reduce_on_oom, + min_chunk_size=min_chunk_size, + clear_cache_on_retry=clear_cache_on_retry, + clear_cache_between_chunks=clear_cache_between_chunks, + ), + ) + decoded = decode_latent_batch_chunked( + sampled, + vae, + decode_chunk_size, + clear_cache_between_chunks=clear_cache_between_chunks, + ) + reset_mask = torch.ones( + (batch_size, images.shape[1], images.shape[2]), + device=images.device, + dtype=images.dtype, + ) + return decoded, sampled, denoised, reset_mask + + previous_input = None + previous_output = None + output_images = [] + output_latents = [] + output_denoised_latents = [] + reset_values = [] + + for frame_index in range(batch_size): + current_frame = _single_frame(images, frame_index) + should_reset = _compute_reset( + frame_index, + previous_input, + current_frame, + int(reset_every_n), + float(scene_cut_threshold), + ) + reset_values.append(float(should_reset)) + + if should_reset: + prior_frame = current_frame + else: + prior_frame, _ = _resolve_temporal_prior( + temporal_mode, + frame_index, + current_frame, + previous_input, + previous_output, + warped_previous_images, + flow_confidence, + float(flow_confidence_scale), + float(temporal_strength), + batch_size, + ) + + encoded = encode_image_batch_chunked( + prior_frame, + vae, + encode_chunk_size, + clear_cache_between_chunks=False, + ) + sampled, denoised = sample_single_latent_temporal( + noise, + guider, + sampler, + sigmas, + encoded, + noise_seed_override=_seed_for_frame(noise, frame_index, bool(lock_seed), seed_stride_mode), + forced_batch_index_value=_batch_index_for_frame(frame_index, bool(lock_seed)), + clear_cache_between_chunks=clear_cache_between_chunks, + ) + decoded = decode_latent_batch_chunked( + sampled, + vae, + decode_chunk_size, + clear_cache_between_chunks=False, + ) + + output_images.append(decoded) + output_latents.append(sampled["samples"]) + output_denoised_latents.append(denoised["samples"]) + + previous_input = current_frame + previous_output = decoded + + if clear_cache_between_chunks and frame_index + 1 < batch_size: + soft_clear_cache() + + images_out = torch.cat(output_images, dim=0) + latents_out = {"samples": torch.cat(output_latents, dim=0)} + denoised_out = {"samples": torch.cat(output_denoised_latents, dim=0)} + reset_mask = torch.tensor(reset_values, device=images.device, dtype=images.dtype).view(batch_size, 1, 1) + reset_mask = reset_mask.expand(batch_size, images.shape[1], images.shape[2]) + return images_out, latents_out, denoised_out, reset_mask + + +NODE_CLASS_MAPPINGS = { + "FluxVideoCleanupTemporalAdvanced": FluxVideoCleanupTemporalAdvanced, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxVideoCleanupTemporalAdvanced": "Flux Video Cleanup Temporal Advanced", +}