Add temporal cleanup node and refactor chunked sampling
This commit is contained in:
xmarre
@@ -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
|
||||
|
||||
|
||||
+12
-1
@@ -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"]
|
||||
|
||||
+15
-331
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
}
|
||||
@@ -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",
|
||||
}
|
||||
Reference in New Issue
Block a user