Add chunked batch nodes and gitignore
This commit is contained in:
xmarre
+30
@@ -0,0 +1,30 @@
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*.pyo
|
||||
*.pyd
|
||||
|
||||
.coverage
|
||||
.coverage.*
|
||||
htmlcov/
|
||||
.mypy_cache/
|
||||
.nox/
|
||||
.pytest_cache/
|
||||
.ruff_cache/
|
||||
.tox/
|
||||
|
||||
.venv/
|
||||
venv/
|
||||
env/
|
||||
|
||||
build/
|
||||
dist/
|
||||
*.egg-info/
|
||||
|
||||
.idea/
|
||||
.vscode/
|
||||
.vs/
|
||||
.codex
|
||||
.codex/
|
||||
@@ -0,0 +1,34 @@
|
||||
# ComfyUI Chunked Batch Nodes
|
||||
|
||||
Custom nodes for chunked batched image-to-image workflows.
|
||||
|
||||
Included nodes:
|
||||
|
||||
- **VAE Encode Batch Chunked**
|
||||
- **SamplerCustomAdvanced Chunked**
|
||||
- **VAE Decode Batch Chunked**
|
||||
|
||||
## 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.
|
||||
|
||||
Typical workflow:
|
||||
|
||||
`IMAGE batch -> VAE Encode Batch Chunked -> SamplerCustomAdvanced Chunked -> VAE Decode Batch Chunked`
|
||||
|
||||
## Notes
|
||||
|
||||
- `SamplerCustomAdvanced Chunked` 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.
|
||||
- `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`.
|
||||
|
||||
## 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.
|
||||
|
||||
## Installation
|
||||
|
||||
Copy this folder into your ComfyUI `custom_nodes` directory and restart ComfyUI.
|
||||
@@ -0,0 +1,3 @@
|
||||
from .chunked_batch_nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,324 @@
|
||||
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:
|
||||
reps = (math.ceil(full_batch_size / value.shape[0]),) + ((1,) * (value.ndim - 1))
|
||||
expanded = value.repeat(reps)[:full_batch_size]
|
||||
return expanded[start:end]
|
||||
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)
|
||||
|
||||
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)
|
||||
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",
|
||||
}
|
||||
Reference in New Issue
Block a user