From 0c4ae6c3297b823d7d5f80b6654b176cff4e5875 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Sun, 19 Jul 2026 19:54:21 +0900 Subject: [PATCH] Add continuous VAE decoding --- comfy/continuous_batching.py | 195 +++++++++++++ nodes.py | 36 +++ .../comfy_test/test_continuous_batching.py | 257 +++++++++++++++++- 3 files changed, 487 insertions(+), 1 deletion(-) diff --git a/comfy/continuous_batching.py b/comfy/continuous_batching.py index bde1041dd..072f02d06 100644 --- a/comfy/continuous_batching.py +++ b/comfy/continuous_batching.py @@ -1,4 +1,5 @@ import asyncio +import contextvars import logging import math from dataclasses import dataclass, field @@ -14,6 +15,7 @@ import comfy.model_patcher import comfy.patcher_extension import comfy.sampler_helpers import comfy.samplers +import comfy.sd import comfy.supported_models from comfy_execution.progress import reset_progress_registry, set_progress_registry from comfy_execution.utils import CurrentNodeContext, reset_current_client_id, set_current_client_id @@ -29,6 +31,7 @@ CONTINUOUS_SAMPLER_NODE_FAMILIES = { } _COORDINATORS = {} +_VAE_COORDINATORS = {} _CANCEL_CHECKER = None @@ -660,3 +663,195 @@ async def sample_euler_continuous(state): coordinator = ContinuousBatchCoordinator(model_key, state) _COORDINATORS[model_key] = coordinator return await coordinator.submit(state) + + +@dataclass +class ContinuousVAEDecodeRequest: + vae: Any + samples: torch.Tensor + max_batch_size: int + admission_delay: float + prompt_id: str | None = None + + def validate(self): + if type(self.vae) is not comfy.sd.VAE: + raise ValueError("Continuous VAE decoding requires a core VAE") + if not torch.is_tensor(self.samples): + raise ValueError("Continuous VAE decoding requires a tensor latent") + if self.samples.is_nested: + raise ValueError("Continuous VAE decoding does not support nested latents") + if self.samples.layout is not torch.strided: + raise ValueError("Continuous VAE decoding requires a dense strided latent") + if self.samples.ndim == 0: + raise ValueError("Continuous VAE decoding requires a batched latent") + if self.samples.shape[0] != 1: + raise ValueError("Continuous VAE decoding requires one latent per request") + if self.vae.latent_dim == 2: + if self.samples.ndim != 4: + raise ValueError("Continuous VAE decoding requires a four-dimensional image latent") + elif self.vae.latent_dim == 3: + if self.samples.ndim != 5 or self.samples.shape[2] != 1: + raise ValueError("Continuous VAE decoding supports only single-frame image latents for three-dimensional VAEs") + else: + raise ValueError("Continuous VAE decoding does not support this VAE latent dimensionality") + if self.max_batch_size < 1: + raise ValueError("Continuous VAE decoding max batch size must be positive") + if not math.isfinite(self.admission_delay) or self.admission_delay < 0: + raise ValueError("Continuous VAE decoding admission delay must be finite and non-negative") + + def key(self): + self.validate() + return ( + id(self.vae), + id(self.vae.patcher), + self.vae.device, + self.vae.vae_dtype, + self.vae.output_device, + tuple(self.samples.shape[1:]), + self.samples.dtype, + self.samples.device, + self.max_batch_size, + self.admission_delay, + ) + + def vae_key(self): + return id(self.vae) + + def clear(self): + self.vae = None + self.samples = None + + +@dataclass +class _QueuedVAERequest: + state: ContinuousVAEDecodeRequest + future: asyncio.Future = field(init=False) + cancelled: bool = False + + def __post_init__(self): + self.future = asyncio.get_running_loop().create_future() + + +class ContinuousVAECoordinator: + def __init__(self, vae_key, vae): + self.vae_key = vae_key + self.vae = vae + self.pending = [] + self.task = None + self.last_batch_size = None + + async def submit(self, state): + request = _QueuedVAERequest(state) + self.pending.append(request) + if self.task is None: + self.task = contextvars.Context().run(asyncio.create_task, self._run()) + try: + return await request.future + except asyncio.CancelledError: + request.cancelled = True + raise + + @staticmethod + def _cancel(request): + request.state.clear() + if not request.future.done(): + request.future.set_exception(comfy.model_management.InterruptProcessingException()) + + @staticmethod + def _fail(request, error): + request.state.clear() + if not request.future.done(): + request.future.set_exception(error) + + def _drop_cancelled(self): + remaining = [] + for request in self.pending: + if request.cancelled or _is_cancelled(request.state.prompt_id): + self._cancel(request) + else: + remaining.append(request) + self.pending = remaining + + def _take_group(self, group_key, max_batch_size): + group = [] + remaining = [] + for request in self.pending: + if request.cancelled or _is_cancelled(request.state.prompt_id): + self._cancel(request) + elif request.state.key() == group_key and len(group) < max_batch_size: + group.append(request) + else: + remaining.append(request) + self.pending = remaining + return group + + def _decode(self, group): + if len(group) == 1: + images = self.vae.decode(group[0].state.samples) + else: + images = self.vae.decode(torch.cat([request.state.samples for request in group])) + if not torch.is_tensor(images) or images.ndim < 1 or images.shape[0] != len(group): + raise RuntimeError("Continuous VAE decoder returned an invalid batch shape") + return images + + def _finish_group(self, group, images): + for index, request in enumerate(group): + if request.cancelled or _is_cancelled(request.state.prompt_id): + self._cancel(request) + continue + output = images if len(group) == 1 else images[index:index + 1].clone() + request.state.clear() + if not request.future.done(): + request.future.set_result(output) + + async def _run(self): + error = None + try: + while self.pending: + self._drop_cancelled() + if not self.pending: + continue + representative = self.pending[0] + group_key = representative.state.key() + if representative.state.admission_delay > 0: + await asyncio.sleep(representative.state.admission_delay) + self._drop_cancelled() + if not any(request is representative for request in self.pending): + continue + group = self._take_group(group_key, representative.state.max_batch_size) + if not group: + continue + if len(group) != self.last_batch_size: + logging.info("Continuous VAE batch size: %d", len(group)) + self.last_batch_size = len(group) + try: + images = self._decode(group) + except Exception as decode_error: + for request in group: + self._fail(request, decode_error) + continue + self._finish_group(group, images) + except Exception as run_error: + error = run_error + finally: + if error is not None: + for request in self.pending: + self._fail(request, error) + else: + for request in self.pending: + self._cancel(request) + self.pending.clear() + self.vae = None + if _VAE_COORDINATORS.get(self.vae_key) is self: + del _VAE_COORDINATORS[self.vae_key] + self.task = None + + +async def decode_vae_continuous(state): + state.key() + vae_key = state.vae_key() + coordinator = _VAE_COORDINATORS.get(vae_key) + if coordinator is None: + coordinator = ContinuousVAECoordinator(vae_key, state.vae) + _VAE_COORDINATORS[vae_key] = coordinator + return await coordinator.submit(state) diff --git a/nodes.py b/nodes.py index f5c45db20..f7b573081 100644 --- a/nodes.py +++ b/nodes.py @@ -340,6 +340,40 @@ class VAEDecode: images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) return (images, ) + +class ContinuousVAEDecode: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "samples": ("LATENT", {"tooltip": "The latent to be decoded."}), + "vae": ("VAE", {"tooltip": "The VAE model used for decoding the latent."}), + "max_batch_size": ("INT", {"default": 4, "min": 1, "max": 64}), + "admission_delay_ms": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 1000.0, "step": 0.1}), + } + } + RETURN_TYPES = ("IMAGE",) + OUTPUT_TOOLTIPS = ("The decoded image.",) + FUNCTION = "decode" + + CATEGORY = "model/latent" + DESCRIPTION = "Continuously batches compatible core VAE decodes from independent workflows." + + async def decode(self, samples, vae, max_batch_size=4, admission_delay_ms=2.0): + execution_context = get_executing_context() + state = comfy.continuous_batching.ContinuousVAEDecodeRequest( + vae=vae, + samples=samples["samples"], + max_batch_size=max_batch_size, + admission_delay=admission_delay_ms / 1000.0, + prompt_id=execution_context.prompt_id if execution_context is not None else None, + ) + images = await comfy.continuous_batching.decode_vae_continuous(state) + if len(images.shape) == 5: #Combine batches + images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1]) + return (images, ) + + class VAEDecodeTiled: @classmethod def INPUT_TYPES(s): @@ -2139,6 +2173,7 @@ NODE_CLASS_MAPPINGS = { "CLIPTextEncode": CLIPTextEncode, "CLIPSetLastLayer": CLIPSetLastLayer, "VAEDecode": VAEDecode, + "ContinuousVAEDecode": ContinuousVAEDecode, "VAEEncode": VAEEncode, "VAEEncodeForInpaint": VAEEncodeForInpaint, "VAELoader": VAELoader, @@ -2251,6 +2286,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VAEEncodeForInpaint": "VAE Encode (for Inpainting)", "SetLatentNoiseMask": "Set Latent Noise Mask", "VAEDecode": "VAE Decode", + "ContinuousVAEDecode": "VAE Decode (Continuous)", "VAEEncode": "VAE Encode", "LatentRotate": "Rotate Latent", "LatentFlip": "Flip Latent", diff --git a/tests-unit/comfy_test/test_continuous_batching.py b/tests-unit/comfy_test/test_continuous_batching.py index 5ca15ef0f..a5b65bfbe 100644 --- a/tests-unit/comfy_test/test_continuous_batching.py +++ b/tests-unit/comfy_test/test_continuous_batching.py @@ -5,6 +5,7 @@ import pytest import torch import comfy.conds +import comfy.continuous_batching as continuous_batching import comfy.model_base import comfy.patcher_extension import comfy.supported_models @@ -14,6 +15,7 @@ from comfy.continuous_batching import ( FAMILY_SDXL, ContinuousBatchCoordinator, ContinuousBatchSession, + ContinuousVAEDecodeRequest, _cfg_branches, _conditioning_structure, _PreparedConditioning, @@ -22,11 +24,12 @@ from comfy.continuous_batching import ( _validate_model_extensions, _validate_model_family, cfg_combine, + decode_vae_continuous, euler_step, ) from comfy_execution.graph import DynamicPrompt from comfy_execution.progress import ProgressRegistry, get_progress_state, reset_progress_state -from comfy_execution.utils import get_current_client_id, get_executing_context, reset_current_client_id, set_current_client_id +from comfy_execution.utils import CurrentNodeContext, get_current_client_id, get_executing_context, reset_current_client_id, set_current_client_id class FakeState: @@ -80,6 +83,38 @@ class FakeSession: self.reload_requests += 1 +class FakeVAE: + def __init__(self, latent_dim=2, not_video=False, fail_shape=None, on_decode=None): + self.patcher = object() + self.device = torch.device("cpu") + self.vae_dtype = torch.float32 + self.output_device = torch.device("cpu") + self.latent_dim = latent_dim + self.not_video = not_video + self.fail_shape = fail_shape + self.on_decode = on_decode + self.calls = [] + + def decode(self, samples): + self.calls.append(samples) + if self.on_decode is not None: + self.on_decode(samples) + if self.fail_shape is not None and tuple(samples.shape[2:]) == self.fail_shape: + self.fail_shape = None + raise RuntimeError("VAE decode failed") + return samples + + +def _vae_request(vae, samples, max_batch_size=4, admission_delay=0.01, prompt_id=None): + return ContinuousVAEDecodeRequest( + vae=vae, + samples=samples, + max_batch_size=max_batch_size, + admission_delay=admission_delay, + prompt_id=prompt_id, + ) + + def _model_patcher(model): return SimpleNamespace(model=model) @@ -486,3 +521,223 @@ def test_failure_clears_requests_and_finishing_request_reloads_survivor(): asyncio.run(failure()) asyncio.run(residency()) + + +def test_continuous_vae_single_request_passes_the_original_tensor(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + + async def run(): + vae = FakeVAE() + latent = torch.full((1, 4, 2, 2), 3.0) + state = _vae_request(vae, latent) + result = await decode_vae_continuous(state) + assert vae.calls == [latent] + assert result is latent + assert state.vae is None + assert state.samples is None + assert not continuous_batching._VAE_COORDINATORS + + asyncio.run(run()) + + +def test_continuous_vae_coalesces_four_requests_in_order_and_clones_outputs(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + + async def run(): + vae = FakeVAE() + latents = [torch.full((1, 4, 2, 2), value, dtype=torch.float32) for value in range(4)] + states = [_vae_request(vae, latent) for latent in latents] + results = await asyncio.gather(*[decode_vae_continuous(state) for state in states]) + assert len(vae.calls) == 1 + assert vae.calls[0].shape == (4, 4, 2, 2) + assert torch.equal(vae.calls[0][:, 0, 0, 0], torch.arange(4, dtype=torch.float32)) + for index, result in enumerate(results): + assert result.shape == (1, 4, 2, 2) + assert torch.equal(result, latents[index]) + assert result.data_ptr() != vae.calls[0][index:index + 1].data_ptr() + assert states[index].vae is None + assert states[index].samples is None + assert not continuous_batching._VAE_COORDINATORS + + asyncio.run(run()) + + +def test_continuous_vae_separates_shapes_dtypes_and_vae_identities(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + + async def run(): + vae = FakeVAE() + other_vae = FakeVAE() + requests = [ + decode_vae_continuous(_vae_request(vae, torch.zeros(1, 4, 2, 2))), + decode_vae_continuous(_vae_request(vae, torch.zeros(1, 4, 3, 2))), + decode_vae_continuous(_vae_request(vae, torch.zeros(1, 4, 2, 2, dtype=torch.float16))), + decode_vae_continuous(_vae_request(vae, torch.zeros(1, 4, 2, 2, device="meta"))), + decode_vae_continuous(_vae_request(other_vae, torch.zeros(1, 4, 2, 2))), + ] + await asyncio.gather(*requests) + assert len(vae.calls) == 4 + assert all(call.shape[0] == 1 for call in vae.calls) + assert len(other_vae.calls) == 1 + assert other_vae.calls[0].shape == (1, 4, 2, 2) + assert not continuous_batching._VAE_COORDINATORS + + asyncio.run(run()) + + +def test_continuous_vae_accepts_and_coalesces_single_frame_anima_latents(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + + async def run(): + vae = FakeVAE(latent_dim=3, not_video=False) + latents = [torch.full((1, 16, 1, 2, 2), value) for value in (1.0, 2.0)] + results = await asyncio.gather(*[decode_vae_continuous(_vae_request(vae, latent)) for latent in latents]) + assert len(vae.calls) == 1 + assert vae.calls[0].shape == (2, 16, 1, 2, 2) + assert all(torch.equal(result, latent) for result, latent in zip(results, latents)) + + asyncio.run(run()) + + +def test_continuous_vae_rejects_unsupported_inputs(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + vae = FakeVAE() + + with pytest.raises(ValueError, match="core VAE"): + _vae_request(object(), torch.zeros(1, 4, 2, 2)).key() + with pytest.raises(ValueError, match="batched latent"): + _vae_request(vae, torch.tensor(0.0)).key() + with pytest.raises(ValueError, match="one latent"): + _vae_request(vae, torch.zeros(2, 4, 2, 2)).key() + with pytest.raises(ValueError, match="four-dimensional"): + _vae_request(vae, torch.zeros(1, 4, 1, 2, 2)).key() + with pytest.raises(ValueError, match="nested"): + _vae_request(vae, torch.nested.nested_tensor([torch.zeros(4, 2, 2)], layout=torch.jagged)).key() + with pytest.raises(ValueError, match="dense strided"): + _vae_request(vae, torch.zeros(1, 4, 2, 2).to_sparse()).key() + with pytest.raises(ValueError, match="single-frame"): + _vae_request(FakeVAE(latent_dim=3, not_video=True), torch.zeros(1, 4, 2, 2, 2)).key() + with pytest.raises(ValueError, match="max batch"): + _vae_request(vae, torch.zeros(1, 4, 2, 2), max_batch_size=0).key() + + +def test_continuous_vae_cancel_after_decode_does_not_cancel_peer(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + cancelled = set() + continuous_batching.set_cancel_checker(cancelled.__contains__) + + async def run(): + vae = FakeVAE(on_decode=lambda samples: cancelled.add("cancelled")) + cancelled_state = _vae_request(vae, torch.zeros(1, 4, 2, 2), prompt_id="cancelled") + peer_state = _vae_request(vae, torch.ones(1, 4, 2, 2), prompt_id="peer") + results = await asyncio.gather( + decode_vae_continuous(cancelled_state), + decode_vae_continuous(peer_state), + return_exceptions=True, + ) + assert isinstance(results[0], comfy.model_management.InterruptProcessingException) + assert torch.equal(results[1], torch.ones(1, 4, 2, 2)) + assert vae.calls[0].shape[0] == 2 + assert cancelled_state.vae is None + assert cancelled_state.samples is None + assert peer_state.vae is None + assert peer_state.samples is None + assert not continuous_batching._VAE_COORDINATORS + + try: + asyncio.run(run()) + finally: + continuous_batching.set_cancel_checker(None) + + +def test_continuous_vae_failure_cleans_group_and_continues_pending_group(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + + async def run(): + vae = FakeVAE(fail_shape=(2, 2)) + failed = torch.zeros(1, 4, 2, 2) + survivor = torch.ones(1, 4, 3, 2) + failed_state = _vae_request(vae, failed) + survivor_state = _vae_request(vae, survivor) + results = await asyncio.gather( + decode_vae_continuous(failed_state), + decode_vae_continuous(survivor_state), + return_exceptions=True, + ) + assert isinstance(results[0], RuntimeError) + assert results[1] is survivor + assert [tuple(call.shape[2:]) for call in vae.calls] == [(2, 2), (3, 2)] + assert failed_state.vae is None + assert failed_state.samples is None + assert survivor_state.vae is None + assert survivor_state.samples is None + assert not continuous_batching._VAE_COORDINATORS + + recovered = torch.full((1, 4, 2, 2), 2.0) + assert await decode_vae_continuous(_vae_request(vae, recovered)) is recovered + assert not continuous_batching._VAE_COORDINATORS + + asyncio.run(run()) + + +def test_continuous_vae_batched_failure_clears_every_request_and_recovers(monkeypatch): + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeVAE) + + async def run(): + vae = FakeVAE(fail_shape=(2, 2)) + states = [_vae_request(vae, torch.full((1, 4, 2, 2), value)) for value in (1.0, 2.0)] + results = await asyncio.gather(*[decode_vae_continuous(state) for state in states], return_exceptions=True) + assert all(isinstance(result, RuntimeError) for result in results) + assert vae.calls[0].shape == (2, 4, 2, 2) + assert all(state.vae is None and state.samples is None for state in states) + assert not continuous_batching._VAE_COORDINATORS + + recovered = torch.zeros(1, 4, 2, 2) + assert await decode_vae_continuous(_vae_request(vae, recovered)) is recovered + assert len(vae.calls) == 2 + assert not continuous_batching._VAE_COORDINATORS + + asyncio.run(run()) + + +def test_continuous_vae_node_preserves_standard_five_dimensional_image_contract(monkeypatch): + import nodes + + class FakeImageVAE(FakeVAE): + def decode(self, samples): + self.calls.append(samples) + return samples.movedim(1, -1) + + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", FakeImageVAE) + + async def run(): + vae = FakeImageVAE(latent_dim=3, not_video=False) + images, = await nodes.ContinuousVAEDecode().decode( + samples={"samples": torch.zeros(1, 4, 1, 2, 2)}, + vae=vae, + max_batch_size=4, + admission_delay_ms=0.0, + ) + assert images.shape == (1, 2, 2, 4) + assert nodes.NODE_CLASS_MAPPINGS["ContinuousVAEDecode"] is nodes.ContinuousVAEDecode + + asyncio.run(run()) + + +def test_continuous_vae_coordinator_runs_without_the_submitter_context(monkeypatch): + seen = [] + + class ContextVAE(FakeVAE): + def decode(self, samples): + seen.append(get_executing_context()) + return super().decode(samples) + + monkeypatch.setattr(continuous_batching.comfy.sd, "VAE", ContextVAE) + + async def run(): + vae = ContextVAE() + with CurrentNodeContext("request", "decode"): + await decode_vae_continuous(_vae_request(vae, torch.zeros(1, 4, 2, 2))) + + asyncio.run(run()) + assert seen == [None]