Add continuous batching samplers

comfy/cli_args.py

@@ -155,6 +155,7 @@ parser.add_argument("--enable-dynamic-vram", action="store_true", help="Enable d
 parser.add_argument("--fast-disk", action="store_true", help="Prefer disk-backed dynamic loading and offload over unpinned RAM. Can be faster for users with fast NVME disks.")
 
 parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")
+parser.add_argument("--continuous-batching", type=int, default=0, metavar="MAX_PROMPTS", help="Run up to MAX_PROMPTS compatible continuous sampler workflows cooperatively using the legacy ModelPatcher. DynamicVRAM is disabled.")
 
 parser.add_argument("--default-hashing-function", type=str, choices=['md5', 'sha1', 'sha256', 'sha512'], default='sha256', help="Allows you to choose the hash function to use for duplicate filename / contents comparison. Default is sha256.")
 
@@ -254,6 +255,13 @@ else:
 if args.cache_ram is not None and len(args.cache_ram) > 2:
     parser.error("--cache-ram accepts at most two values: active GB and inactive GB")
 
+if args.continuous_batching < 0:
+    parser.error("--continuous-batching must be zero or greater")
+if args.continuous_batching > 1 and args.cache_none:
+    parser.error("--continuous-batching with more than one prompt is incompatible with --cache-none")
+if args.continuous_batching and args.enable_dynamic_vram:
+    parser.error("--continuous-batching is incompatible with --enable-dynamic-vram")
+
 if args.high_ram:
     args.cache_classic = True
 
@@ -282,6 +290,8 @@ else:
     args.fast = set(args.fast)
 
 def enables_dynamic_vram():
+    if args.continuous_batching:
+        return False
     if args.enable_dynamic_vram:
         return True
     return not args.disable_dynamic_vram and not args.highvram and not args.gpu_only and not args.novram and not args.cpu

comfy/continuous_batching.py

@@ -0,0 +1,645 @@
+import asyncio
+import logging
+import math
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+
+import comfy.conds
+import comfy.hooks
+import comfy.model_base
+import comfy.model_management
+import comfy.model_patcher
+import comfy.patcher_extension
+import comfy.sampler_helpers
+import comfy.samplers
+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
+
+
+FAMILY_ANIMA = "anima"
+FAMILY_SD15 = "sd15"
+FAMILY_SDXL = "sdxl"
+CONTINUOUS_SAMPLER_NODE_FAMILIES = {
+    "AnimaContinuousKSampler": FAMILY_ANIMA,
+    "SD15ContinuousKSampler": FAMILY_SD15,
+    "SDXLContinuousKSampler": FAMILY_SDXL,
+}
+
+_COORDINATORS = {}
+_CANCEL_CHECKER = None
+
+
+def set_cancel_checker(checker):
+    global _CANCEL_CHECKER
+    _CANCEL_CHECKER = checker
+
+
+def _is_cancelled(prompt_id):
+    return prompt_id is not None and _CANCEL_CHECKER is not None and _CANCEL_CHECKER(prompt_id)
+
+
+def euler_step(x, denoised, sigma, sigma_next):
+    return x + (x - denoised) / sigma * (sigma_next - sigma)
+
+
+_UNSUPPORTED_CONDITIONING = {
+    "additional_models",
+    "area",
+    "clip_end_percent",
+    "clip_start_percent",
+    "control",
+    "crossattn_controlnet",
+    "default",
+    "end_percent",
+    "gligen",
+    "hooks",
+    "mask",
+    "mask_strength",
+    "noise_concat",
+    "start_percent",
+    "strength",
+    "timestep_end",
+    "timestep_start",
+    "unclip_conditioning",
+}
+
+
+def _validate_conditioning(name, conditions):
+    if len(conditions) != 1:
+        raise ValueError(f"Continuous batching requires one {name} conditioning entry")
+    present = _UNSUPPORTED_CONDITIONING.intersection(conditions[0][1])
+    if present:
+        raise ValueError("Continuous batching does not support conditioning feature(s): {}".format(", ".join(sorted(present))))
+
+
+def _value_structure(value):
+    if torch.is_tensor(value):
+        return ("tensor", tuple(value.shape), value.dtype, value.device)
+    if isinstance(value, dict):
+        return ("dict", tuple((key, _value_structure(item)) for key, item in sorted(value.items())))
+    if isinstance(value, (list, tuple)):
+        return (type(value), tuple(_value_structure(item) for item in value))
+    return type(value)
+
+
+def _conditioning_structure(family, conditions):
+    value, metadata = conditions[0]
+    if family == FAMILY_ANIMA:
+        text_ids = metadata.get("t5xxl_ids")
+        if torch.is_tensor(text_ids):
+            return ("anima_context", max(512, text_ids.shape[-1]), value.shape[-1], value.dtype)
+    return (
+        _value_structure(value),
+        tuple((key, _value_structure(item)) for key, item in sorted(metadata.items())),
+    )
+
+
+def _nested_mapping_entries(mapping):
+    if not isinstance(mapping, dict):
+        return bool(mapping)
+    return any(_nested_mapping_entries(value) for value in mapping.values())
+
+
+def _wrapper_types(wrappers):
+    present = set()
+    for wrapper_type, keyed_wrappers in wrappers.items():
+        if _nested_mapping_entries(keyed_wrappers):
+            present.add(getattr(wrapper_type, "value", wrapper_type))
+    return present
+
+
+def _model_wrapper_types(model_patcher):
+    transformer_options = model_patcher.model_options.get("transformer_options", {})
+    wrappers = {}
+    comfy.patcher_extension.merge_nested_dicts(wrappers, model_patcher.wrappers, copy_dict1=False)
+    comfy.patcher_extension.merge_nested_dicts(wrappers, transformer_options.get("wrappers", {}), copy_dict1=False)
+    return _wrapper_types(wrappers)
+
+
+def _validate_model_extensions(family, model_patcher):
+    model_options = model_patcher.model_options
+    unsupported_options = {
+        "context_handler",
+        "model_function_wrapper",
+        "multigpu_clones",
+        "sampler_calc_cond_batch_function",
+        "sampler_cfg_function",
+        "sampler_post_cfg_function",
+        "sampler_pre_cfg_function",
+    }
+    present = unsupported_options.intersection(model_options)
+    if present:
+        raise ValueError("Continuous batching does not support model option(s): {}".format(", ".join(sorted(present))))
+
+    if _nested_mapping_entries(model_patcher.callbacks):
+        raise ValueError("Continuous batching does not support model callbacks")
+    if any(model_patcher.additional_models.values()):
+        raise ValueError("Continuous batching does not support additional models")
+
+    transformer_options = model_options.get("transformer_options", {})
+    if _nested_mapping_entries(transformer_options.get("callbacks", {})):
+        raise ValueError("Continuous batching does not support transformer callbacks")
+
+    present_wrappers = _model_wrapper_types(model_patcher)
+    if present_wrappers:
+        raise ValueError("Continuous batching does not support wrapper(s): {}".format(", ".join(sorted(present_wrappers))))
+
+    patch_keys = {"patches", "patches_replace"}.intersection(transformer_options)
+    if any(_nested_mapping_entries(transformer_options[key]) for key in patch_keys):
+        raise ValueError("Continuous batching does not support transformer patches")
+
+
+def _validate_model_family(family, model_patcher):
+    model = model_patcher.model
+    if family == FAMILY_ANIMA:
+        if not isinstance(model, comfy.model_base.Anima):
+            raise ValueError("Anima continuous batching requires an Anima model")
+        return
+
+    latent_channels = model.latent_format.latent_channels
+    input_channels = getattr(model.diffusion_model, "in_channels", None)
+    if family == FAMILY_SD15:
+        if type(model) is not comfy.model_base.BaseModel or not isinstance(model.model_config, comfy.supported_models.SD15):
+            raise ValueError("SD1.5 continuous batching requires a standard SD1.5 model")
+    elif family == FAMILY_SDXL:
+        if type(model) not in (comfy.model_base.SDXL, comfy.model_base.SDXLRefiner):
+            raise ValueError("SDXL continuous batching requires a standard SDXL or SDXL Refiner model")
+    else:
+        raise ValueError(f"Unknown continuous batching model family: {family}")
+
+    if input_channels != latent_channels or model.concat_keys:
+        raise ValueError("SD continuous batching does not support inpaint or concatenated-input models")
+
+
+def _processed_conditioning_signature(family, cond):
+    expected = {
+        FAMILY_ANIMA: {"c_crossattn": comfy.conds.CONDRegular},
+        FAMILY_SD15: {"c_crossattn": comfy.conds.CONDCrossAttn},
+        FAMILY_SDXL: {"c_crossattn": comfy.conds.CONDCrossAttn, "y": comfy.conds.CONDRegular},
+    }[family]
+    if cond is None or cond.area is not None or cond.control is not None or cond.patches is not None or cond.hooks is not None:
+        raise ValueError("Continuous batching received unsupported processed conditioning")
+    if cond.conditioning.keys() != expected.keys():
+        present = ", ".join(sorted(cond.conditioning))
+        raise ValueError(f"Continuous batching received unsupported processed conditioning keys: {present}")
+
+    signature = []
+    for key, expected_type in expected.items():
+        value = cond.conditioning[key]
+        if type(value) is not expected_type or not torch.is_tensor(value.cond):
+            raise ValueError(f"Continuous batching received an unsupported {key} conditioning wrapper")
+        signature.append((key, type(value), tuple(value.cond.shape), value.cond.dtype, value.cond.device))
+    return tuple(signature)
+
+
+def cfg_combine(cond, uncond, cfg):
+    if math.isclose(cfg, 1.0):
+        return cond
+    return uncond + (cond - uncond) * cfg
+
+
+def _cfg_branches(cfg, model_options):
+    if math.isclose(cfg, 1.0) and not model_options.get("disable_cfg1_optimization", False):
+        return (("positive", 0),)
+    return (("negative", 1), ("positive", 0))
+
+
+@dataclass
+class ContinuousBatchRequest:
+    family: str
+    model_patcher: Any
+    noise: torch.Tensor
+    latent_image: torch.Tensor
+    positive: list
+    negative: list
+    sigmas: torch.Tensor
+    callback: Any
+    seed: int
+    cfg: float
+    max_batch_size: int
+    admission_delay: float
+    prompt_id: str | None = None
+    node_id: str | None = None
+    client_id: str | None = None
+    progress_registry: Any = None
+    index: int = 0
+    x: torch.Tensor | None = None
+    conds: dict | None = None
+    output: torch.Tensor | None = None
+    prepared: bool = False
+
+    def validate(self):
+        _validate_model_family(self.family, self.model_patcher)
+        if self.model_patcher.is_dynamic():
+            raise ValueError("Continuous batching does not support dynamic model patchers")
+        if self.latent_image.is_nested or self.noise.is_nested:
+            raise ValueError("Continuous batching does not support nested latents")
+        if self.latent_image.shape != self.noise.shape or self.latent_image.shape[0] != 1:
+            raise ValueError("Continuous batching requires one latent per request")
+        if self.sigmas.ndim != 1 or len(self.sigmas) < 2:
+            raise ValueError("Continuous batching requires a one-dimensional sigma schedule")
+        if self.max_batch_size < 1:
+            raise ValueError("Continuous batching max batch size must be positive")
+        if not math.isfinite(self.cfg) or self.cfg < 0:
+            raise ValueError("Continuous batching CFG must be finite and non-negative")
+        _validate_conditioning("positive", self.positive)
+        _validate_conditioning("negative", self.negative)
+        _validate_model_extensions(self.family, self.model_patcher)
+
+    def key(self):
+        self.validate()
+        return (
+            self.family,
+            self.model_key(),
+            tuple(self.latent_image.shape[1:]),
+            _conditioning_structure(self.family, self.positive),
+            _conditioning_structure(self.family, self.negative),
+            self.max_batch_size,
+            self.admission_delay,
+        )
+
+    def model_key(self):
+        return (
+            id(self.model_patcher.model),
+            id(self.model_patcher),
+            self.model_patcher.patches_uuid,
+            self.model_patcher.load_device,
+            self.model_patcher.model_dtype(),
+        )
+
+    def clear(self):
+        self.model_patcher = None
+        self.noise = None
+        self.latent_image = None
+        self.positive = None
+        self.negative = None
+        self.sigmas = None
+        self.callback = None
+        self.progress_registry = None
+        self.x = None
+        self.conds = None
+        self.output = None
+
+
+@dataclass
+class _QueuedRequest:
+    state: ContinuousBatchRequest
+    future: asyncio.Future = field(init=False)
+    cancelled: bool = False
+
+    def __post_init__(self):
+        self.future = asyncio.get_running_loop().create_future()
+
+
+class ContinuousBatchSession:
+    def __init__(self, model_patcher):
+        self.model_patcher = model_patcher
+        self.model_options = None
+        self.inner_model = None
+        self.loaded_models = None
+        self.load_conds = None
+        self.capacity = 0
+        self.orig_hook_mode = None
+        self.reload_model = False
+        self.open = False
+
+    def open_session(self, states):
+        if self.open:
+            return
+        for state in states:
+            state.validate()
+        representative = states[0]
+        combined_conds = {
+            "positive": comfy.sampler_helpers.convert_cond(representative.positive),
+            "negative": comfy.sampler_helpers.convert_cond(representative.negative),
+        }
+        self.model_options = comfy.model_patcher.create_model_options_clone(self.model_patcher.model_options)
+        self.model_options.setdefault("transformer_options", {})["sample_sigmas"] = representative.sigmas
+        comfy.samplers.preprocess_conds_hooks(combined_conds)
+        self.orig_hook_mode = self.model_patcher.hook_mode
+        self.model_patcher.hook_mode = comfy.hooks.EnumHookMode.MinVram
+        self.open = True
+        try:
+            comfy.sampler_helpers.prepare_model_patcher(self.model_patcher, combined_conds, self.model_options)
+            comfy.samplers.filter_registered_hooks_on_conds(combined_conds, self.model_options)
+            self.inner_model, _, self.loaded_models = comfy.sampler_helpers.prepare_sampling(
+                self.model_patcher,
+                [len(states)] + list(representative.noise.shape[1:]),
+                combined_conds,
+                self.model_options,
+            )
+            self.load_conds = combined_conds
+            self.capacity = len(states)
+            comfy.samplers.cast_to_load_options(self.model_options, device=self.model_patcher.load_device, dtype=self.model_patcher.model_dtype())
+            self.model_patcher.pre_run()
+            self.reload_model = False
+        except BaseException:
+            self.close()
+            raise
+
+    def request_model_reload(self):
+        if self.open:
+            self.reload_model = True
+
+    def ensure_model_loaded(self, states):
+        required_capacity = max(self.capacity, len(states))
+        if not self.reload_model and required_capacity == self.capacity:
+            return
+        representative = states[0]
+        self.inner_model, _, self.loaded_models = comfy.sampler_helpers.prepare_sampling(
+            self.model_patcher,
+            [required_capacity] + list(representative.noise.shape[1:]),
+            self.load_conds,
+            self.model_options,
+        )
+        comfy.samplers.cast_to_load_options(self.model_options, device=self.model_patcher.load_device, dtype=self.model_patcher.model_dtype())
+        self.capacity = required_capacity
+        self.reload_model = False
+
+    def prepare_request(self, state):
+        if state.prepared:
+            return
+        device = self.model_patcher.load_device
+        state.noise = state.noise.to(device=device, dtype=torch.float32)
+        state.latent_image = state.latent_image.to(device=device, dtype=torch.float32)
+        state.sigmas = state.sigmas.to(device)
+        conds = {
+            "positive": comfy.sampler_helpers.convert_cond(state.positive),
+            "negative": comfy.sampler_helpers.convert_cond(state.negative),
+        }
+        conds = comfy.samplers.process_conds(
+            self.inner_model,
+            state.noise,
+            conds,
+            device,
+            state.latent_image,
+            None,
+            state.seed,
+            latent_shapes=[state.latent_image.shape],
+        )
+        latent_image = state.latent_image
+        if torch.count_nonzero(latent_image) > 0:
+            latent_image = self.inner_model.process_latent_in(latent_image)
+        sigma = float(state.sigmas[0])
+        sigma_max = float(self.inner_model.model_sampling.sigma_max)
+        max_denoise = math.isclose(sigma_max, sigma, rel_tol=1e-5) or sigma > sigma_max
+        state.x = self.inner_model.model_sampling.noise_scaling(state.sigmas[0], state.noise, latent_image, max_denoise)
+        sigma = state.sigmas[0].to(state.x).unsqueeze(0)
+        for name in ("positive", "negative"):
+            if len(conds[name]) != 1:
+                raise ValueError(f"Continuous batching requires one processed {name} conditioning entry")
+            processed = comfy.samplers.get_area_and_mult(conds[name][0], state.x, sigma)
+            _processed_conditioning_signature(state.family, processed)
+        state.latent_image = latent_image
+        state.conds = conds
+        state.prepared = True
+
+    def predict(self, states):
+        if len(states) == 1:
+            state = states[0]
+            sigma = state.sigmas[state.index].to(state.x).unsqueeze(0)
+            model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
+            model_options.setdefault("transformer_options", {})["sample_sigmas"] = state.sigmas
+            return [comfy.samplers.sampling_function(
+                self.inner_model,
+                state.x,
+                sigma,
+                state.conds["negative"],
+                state.conds["positive"],
+                state.cfg,
+                model_options=model_options,
+                seed=state.seed,
+            )]
+
+        state_timesteps = []
+        state_branches = []
+        entries = []
+        for state_index, state in enumerate(states):
+            sigma = state.sigmas[state.index].to(state.x)
+            state_timesteps.append(sigma)
+            branches = _cfg_branches(state.cfg, self.model_options)
+            state_branches.append(branches)
+            for name, branch in branches:
+                cond = comfy.samplers.get_area_and_mult(state.conds[name][0], state.x, sigma.unsqueeze(0))
+                signature = _processed_conditioning_signature(state.family, cond)
+                entries.append((state_index, name, branch, state.x, sigma, cond, signature))
+
+        buckets = []
+        for entry in entries:
+            for bucket in buckets:
+                if entry[6] == bucket[0][6] and comfy.samplers.can_concat_cond(entry[5], bucket[0][5]):
+                    bucket.append(entry)
+                    break
+            else:
+                buckets.append([entry])
+
+        state_timestep = torch.stack(state_timesteps)
+        self.model_patcher.prepare_state(state_timestep, self.model_options)
+        branch_outputs = [{} for _ in states]
+        for bucket in buckets:
+            input_x = torch.cat([entry[3] for entry in bucket])
+            timestep = torch.stack([entry[4] for entry in bucket])
+            cond_objects = [entry[5] for entry in bucket]
+            conditioning = comfy.samplers.cond_cat([cond.conditioning for cond in cond_objects])
+            transformer_options = self.model_patcher.apply_hooks(hooks=None)
+            if "transformer_options" in self.model_options:
+                transformer_options = comfy.patcher_extension.merge_nested_dicts(
+                    transformer_options,
+                    self.model_options["transformer_options"],
+                    copy_dict1=False,
+                )
+            transformer_options["cond_or_uncond"] = [entry[2] for entry in bucket]
+            transformer_options["uuids"] = [cond.uuid for cond in cond_objects]
+            transformer_options["sigmas"] = timestep
+            conditioning["transformer_options"] = transformer_options
+            outputs = self.inner_model.apply_model(input_x, timestep, **conditioning).split(1)
+            for entry, output in zip(bucket, outputs):
+                branch_outputs[entry[0]][entry[1]] = output
+
+        predictions = []
+        for state, branches, outputs in zip(states, state_branches, branch_outputs):
+            if len(branches) == 1:
+                predictions.append(outputs["positive"])
+            else:
+                predictions.append(cfg_combine(outputs["positive"], outputs["negative"], state.cfg))
+        return predictions
+
+    @staticmethod
+    def run_callback(state, prediction):
+        if state.callback is None:
+            return
+        client_token = set_current_client_id(state.client_id)
+        progress_token = set_progress_registry(state.progress_registry) if state.progress_registry is not None else None
+        try:
+            if state.prompt_id is not None and state.node_id is not None:
+                with CurrentNodeContext(state.prompt_id, state.node_id):
+                    state.callback(state.index, prediction, state.x, len(state.sigmas) - 1)
+            else:
+                state.callback(state.index, prediction, state.x, len(state.sigmas) - 1)
+        finally:
+            if progress_token is not None:
+                reset_progress_registry(progress_token)
+            reset_current_client_id(client_token)
+
+    def step(self, states):
+        with comfy.model_management.cuda_device_context(self.model_patcher.load_device):
+            self.open_session(states)
+            self.ensure_model_loaded(states)
+            for state in states:
+                self.prepare_request(state)
+            denoised = self.predict(states)
+            updates = []
+            for state, prediction in zip(states, denoised):
+                if prediction.shape != state.x.shape:
+                    raise RuntimeError("Continuous batch denoiser returned an invalid shape")
+                sigma = state.sigmas[state.index].to(state.x)
+                self.run_callback(state, prediction)
+                state.x = euler_step(state.x, prediction, sigma, state.sigmas[state.index + 1].to(state.x))
+                state.index += 1
+                finished = state.index == len(state.sigmas) - 1
+                if finished:
+                    samples = self.inner_model.model_sampling.inverse_noise_scaling(state.sigmas[-1], state.x)
+                    samples = self.inner_model.process_latent_out(samples.to(torch.float32))
+                    state.output = samples.to(
+                        device=comfy.model_management.intermediate_device(),
+                        dtype=comfy.model_management.intermediate_dtype(),
+                    )
+                updates.append((state, finished))
+            return updates
+
+    def close(self):
+        if not self.open:
+            return
+        try:
+            comfy.samplers.cast_to_load_options(self.model_options, device=self.model_patcher.offload_device)
+            self.model_patcher.cleanup()
+            if self.loaded_models is not None:
+                comfy.sampler_helpers.cleanup_models({}, self.loaded_models)
+        finally:
+            self.model_patcher.hook_mode = self.orig_hook_mode
+            self.model_patcher.restore_hook_patches()
+            self.inner_model = None
+            self.loaded_models = None
+            self.load_conds = None
+            self.capacity = 0
+            self.model_options = None
+            self.reload_model = False
+            self.open = False
+
+
+class ContinuousBatchCoordinator:
+    def __init__(self, model_key, state):
+        self.model_key = model_key
+        self.session = ContinuousBatchSession(state.model_patcher)
+        self.group_key = None
+        self.pending = []
+        self.active = []
+        self.task = None
+        self.last_batch_size = None
+
+    async def submit(self, state):
+        request = _QueuedRequest(state)
+        self.pending.append(request)
+        if self.task is None:
+            self.task = asyncio.create_task(self._run())
+        try:
+            return await request.future
+        except asyncio.CancelledError:
+            request.cancelled = True
+            raise
+
+    def _admit(self):
+        if not self.pending:
+            return
+        group_key = self.active[0].state.key() if self.active else self.pending[0].state.key()
+        if not self.active:
+            if self.group_key is not None and self.group_key != group_key:
+                self.session.close()
+                self.session = ContinuousBatchSession(self.pending[0].state.model_patcher)
+            self.group_key = group_key
+        max_batch_size = self.active[0].state.max_batch_size if self.active else self.pending[0].state.max_batch_size
+        remaining = []
+        admitted = False
+        for request in self.pending:
+            if request.state.key() != group_key or len(self.active) >= max_batch_size:
+                remaining.append(request)
+                continue
+            if request.cancelled:
+                request.state.clear()
+            else:
+                self.active.append(request)
+                admitted = True
+        self.pending = remaining
+        if admitted and self.session.open:
+            self.session.request_model_reload()
+
+    def _drop_cancelled(self):
+        kept = []
+        for request in self.active:
+            if request.cancelled or _is_cancelled(getattr(request.state, "prompt_id", None)):
+                request.state.clear()
+                if not request.future.done():
+                    request.future.set_exception(comfy.model_management.InterruptProcessingException())
+            else:
+                kept.append(request)
+        self.active = kept
+
+    async def _run(self):
+        try:
+            if self.pending and self.pending[0].state.admission_delay > 0:
+                await asyncio.sleep(self.pending[0].state.admission_delay)
+            while self.pending or self.active:
+                self._admit()
+                self._drop_cancelled()
+                if not self.active:
+                    await asyncio.sleep(0)
+                    continue
+                current = list(self.active)
+                if len(current) != self.last_batch_size:
+                    logging.info("Continuous %s batch size: %d", current[0].state.family, len(current))
+                    self.last_batch_size = len(current)
+                updates = self.session.step([request.state for request in current])
+                self.active = []
+                finished_any = False
+                for request, (state, finished) in zip(current, updates):
+                    if request.cancelled:
+                        state.clear()
+                    elif finished:
+                        finished_any = True
+                        output = state.output
+                        state.clear()
+                        if not request.future.done():
+                            request.future.set_result(output)
+                    else:
+                        self.active.append(request)
+                if finished_any and self.active:
+                    self.session.request_model_reload()
+                self._admit()
+                self._drop_cancelled()
+                if self.pending or self.active:
+                    await asyncio.sleep(0)
+        except BaseException as error:
+            for request in self.active + self.pending:
+                request.state.clear()
+                if not request.future.done():
+                    request.future.set_exception(error)
+            self.active.clear()
+            self.pending.clear()
+        finally:
+            self.session.close()
+            if _COORDINATORS.get(self.model_key) is self:
+                del _COORDINATORS[self.model_key]
+            self.task = None
+
+
+async def sample_euler_continuous(state):
+    state.key()
+    model_key = state.model_key()
+    coordinator = _COORDINATORS.get(model_key)
+    if coordinator is None:
+        coordinator = ContinuousBatchCoordinator(model_key, state)
+        _COORDINATORS[model_key] = coordinator
+    return await coordinator.submit(state)

comfy/sd1_clip.py

@@ -276,7 +276,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         else:
             intermediate_output = self.layer_idx
 
-        outputs = self.transformer(None, attention_mask_model, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=self.layer_norm_hidden_state, dtype=torch.float32, embeds_info=embeds_info)
+        outputs = self.transformer_forward(tokens, attention_mask_model, embeds, num_tokens, intermediate_output, embeds_info)
 
         if self.layer == "last":
             z = outputs[0].float()
@@ -302,6 +302,9 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
 
         return z, pooled_output
 
+    def transformer_forward(self, tokens, attention_mask, embeds, num_tokens, intermediate_output, embeds_info):
+        return self.transformer(None, attention_mask, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=self.layer_norm_hidden_state, dtype=torch.float32, embeds_info=embeds_info)
+
     def encode(self, tokens):
         return self(tokens)
 

comfy/text_encoders/anima.py

@@ -1,5 +1,6 @@
 from transformers import Qwen2Tokenizer, T5TokenizerFast
 import comfy.text_encoders.llama
+import comfy.text_encoders.anima_cache
 from comfy import sd1_clip
 import os
 import torch
@@ -40,10 +41,14 @@ class Qwen3_06BModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen3_06B, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
+    def transformer_forward(self, tokens, attention_mask, embeds, num_tokens, intermediate_output, embeds_info):
+        return comfy.text_encoders.anima_cache.forward(self.transformer, comfy.text_encoders.anima_cache.get_owner(self), tokens, attention_mask, embeds, num_tokens, intermediate_output, self.layer_norm_hidden_state, torch.float32, embeds_info)
+
 
 class AnimaTEModel(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, model_options={}):
         super().__init__(device=device, dtype=dtype, name="qwen3_06b", clip_model=Qwen3_06BModel, model_options=model_options)
+        comfy.text_encoders.anima_cache.set_owner(self.qwen3_06b, self)
 
     def encode_token_weights(self, token_weight_pairs):
         out = super().encode_token_weights(token_weight_pairs)

comfy/text_encoders/anima_cache.py

@@ -0,0 +1,126 @@
+import contextvars
+import logging
+import numbers
+import weakref
+
+import torch
+
+
+_cache_scope = contextvars.ContextVar("anima_prefix_cache", default=None)
+
+
+def begin_cache_scope(enabled=True):
+    cache = weakref.WeakKeyDictionary() if enabled else None
+    return cache, _cache_scope.set(cache)
+
+
+def end_cache_scope(scope):
+    cache, token = scope
+    try:
+        if cache is not None:
+            cache.clear()
+    finally:
+        _cache_scope.reset(token)
+
+
+def set_owner(model, owner):
+    object.__setattr__(model, "_anima_cache_owner", weakref.ref(owner))
+
+
+def get_owner(model):
+    owner = getattr(model, "_anima_cache_owner", None)
+    return owner() if owner is not None else None
+
+
+def _token_ids(tokens):
+    if len(tokens) != 1:
+        return None
+    sequence = tokens[0]
+    if torch.is_tensor(sequence):
+        sequence = sequence.tolist()
+    token_ids = []
+    for token in sequence:
+        if isinstance(token, numbers.Integral):
+            token_ids.append(int(token))
+        elif isinstance(token, (tuple, list)) and len(token) == 2 and isinstance(token[0], numbers.Integral) and isinstance(token[1], numbers.Real) and token[1] == 1:
+            token_ids.append(int(token[0]))
+        else:
+            return None
+    return tuple(token_ids)
+
+
+def _common_prefix(first, second):
+    length = min(len(first), len(second))
+    for index in range(length):
+        if first[index] != second[index]:
+            return index
+    return length
+
+
+def _copy_key_values(key_values, length):
+    return [(key[:, :, :length].clone(), value[:, :, :length].clone(), length) for key, value, _ in key_values]
+
+
+def forward(transformer, cache_owner, tokens, attention_mask, embeds, num_tokens, intermediate_output, final_layer_norm_intermediate, dtype, embeds_info):
+    prefix_cache = _cache_scope.get()
+    token_ids = _token_ids(tokens)
+    weight_uuid = getattr(cache_owner, "current_weight_patches_uuid", None)
+    if prefix_cache is None or cache_owner is None or weight_uuid is None or token_ids is None or embeds_info or intermediate_output is not None or attention_mask is not None and not bool(torch.all(attention_mask)):
+        return transformer(None, attention_mask, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=final_layer_norm_intermediate, dtype=dtype, embeds_info=embeds_info)
+
+    cached = prefix_cache.get(cache_owner)
+    if cached is not None and (cached[0]() is not transformer or cached[1] != weight_uuid or cached[3].device != embeds.device or cached[3].dtype != embeds.dtype):
+        cached = None
+    common = _common_prefix(cached[2], token_ids) if cached is not None else 0
+    if common > 0:
+        logging.debug("Anima Qwen cache reused %d prefix tokens", common)
+    if cached is not None and common == len(token_ids) == len(cached[2]):
+        return cached[3].clone(), None
+
+    if cached is not None and common == len(token_ids):
+        hidden = cached[3][:, :common].clone()
+        prefix_cache[cache_owner] = (weakref.ref(transformer), weight_uuid, token_ids, hidden.clone(), _copy_key_values(cached[4], common))
+        return hidden, None
+
+    past_key_values = []
+    prefix_hidden = None
+    if cached is not None and common > 0:
+        prefix_hidden = cached[3][:, :common].clone()
+        past_key_values = _copy_key_values(cached[4], common)
+
+    suffix_embeds = embeds[:, common:]
+    if common > 0 and suffix_embeds.shape[1] > 1:
+        suffix_outputs = []
+        next_key_values = past_key_values
+        for index in range(suffix_embeds.shape[1]):
+            output = transformer(
+                None,
+                None,
+                embeds=suffix_embeds[:, index:index + 1],
+                num_tokens=[1],
+                intermediate_output=intermediate_output,
+                final_layer_norm_intermediate=final_layer_norm_intermediate,
+                dtype=dtype,
+                embeds_info=embeds_info,
+                past_key_values=next_key_values,
+            )
+            suffix_outputs.append(output[0])
+            next_key_values = output[2]
+        suffix_hidden = torch.cat(suffix_outputs, dim=1)
+    else:
+        output = transformer(
+            None,
+            attention_mask,
+            embeds=suffix_embeds,
+            num_tokens=[1] if common > 0 else num_tokens,
+            intermediate_output=intermediate_output,
+            final_layer_norm_intermediate=final_layer_norm_intermediate,
+            dtype=dtype,
+            embeds_info=embeds_info,
+            past_key_values=past_key_values,
+        )
+        suffix_hidden = output[0]
+        next_key_values = output[2]
+    hidden = suffix_hidden if prefix_hidden is None else torch.cat((prefix_hidden, suffix_hidden), dim=1)
+    prefix_cache[cache_owner] = (weakref.ref(transformer), weight_uuid, token_ids, hidden.clone(), _copy_key_values(next_key_values, len(token_ids)))
+    return hidden, None

comfy_execution/caching.py

@@ -1,12 +1,12 @@
 import asyncio
 import bisect
+import contextvars
 import itertools
 import psutil
 import time
 import torch
 from typing import Sequence, Mapping, Dict
 from comfy.model_patcher import ModelPatcher
-from comfy_execution.graph import DynamicPrompt
 from abc import ABC, abstractmethod
 
 import nodes
@@ -150,30 +150,66 @@ class CacheKeySetInputSignature(CacheKeySet):
 class BasicCache:
     def __init__(self, key_class, enable_providers=False):
         self.key_class = key_class
-        self.initialized = False
         self.enable_providers = enable_providers
-        self.dynprompt: DynamicPrompt
-        self.cache_key_set: CacheKeySet
+        self._prompt_context = contextvars.ContextVar("cache_prompt_context", default=None)
+        self._active_key_sets = set()
         self.cache = {}
         self.subcaches = {}
         self._pending_store_tasks: set = set()
 
     async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
-        self.dynprompt = dynprompt
-        self.cache_key_set = self.key_class(dynprompt, node_ids, is_changed_cache)
-        await self.cache_key_set.add_keys(node_ids)
-        self.is_changed_cache = is_changed_cache
-        self.initialized = True
+        previous = self._prompt_context.get()
+        if previous is not None:
+            self._active_key_sets.discard(previous[1])
+        cache_key_set = self.key_class(dynprompt, node_ids, is_changed_cache)
+        await cache_key_set.add_keys(node_ids)
+        self._prompt_context.set((dynprompt, cache_key_set, is_changed_cache))
+        self._active_key_sets.add(cache_key_set)
+
+    @property
+    def initialized(self):
+        return self._prompt_context.get() is not None
+
+    @property
+    def dynprompt(self):
+        context = self._prompt_context.get()
+        return context[0] if context is not None else None
+
+    @property
+    def cache_key_set(self):
+        context = self._prompt_context.get()
+        return context[1] if context is not None else None
+
+    @property
+    def is_changed_cache(self):
+        context = self._prompt_context.get()
+        return context[2] if context is not None else None
+
+    def release_prompt(self):
+        context = self._prompt_context.get()
+        if context is None:
+            return
+        self._active_key_sets.discard(context[1])
+        self._prompt_context.set(None)
+        for subcache in self.subcaches.values():
+            subcache.release_prompt()
+
+    def _active_data_keys(self):
+        keys = set()
+        for key_set in self._active_key_sets:
+            keys.update(key_set.get_used_keys())
+        return keys
 
     def all_node_ids(self):
         assert self.initialized
         node_ids = self.cache_key_set.all_node_ids()
         for subcache in self.subcaches.values():
-            node_ids = node_ids.union(subcache.all_node_ids())
+            if subcache.initialized:
+                node_ids = node_ids.union(subcache.all_node_ids())
         return node_ids
 
     def _clean_cache(self):
-        preserve_keys = set(self.cache_key_set.get_used_keys())
+        preserve_keys = self._active_data_keys()
         to_remove = []
         for key in self.cache:
             if key not in preserve_keys:
@@ -182,7 +218,9 @@ class BasicCache:
             del self.cache[key]
 
     def _clean_subcaches(self):
-        preserve_subcaches = set(self.cache_key_set.get_used_subcache_keys())
+        preserve_subcaches = set()
+        for key_set in self._active_key_sets:
+            preserve_subcaches.update(key_set.get_used_subcache_keys())
 
         to_remove = []
         for key in self.subcaches:
@@ -418,6 +456,9 @@ class NullCache:
     def clean_unused(self):
         pass
 
+    def release_prompt(self):
+        pass
+
     def poll(self, **kwargs):
         pass
 
@@ -454,7 +495,8 @@ class LRUCache(BasicCache):
     def clean_unused(self):
         while len(self.cache) > self.max_size and self.min_generation < self.generation:
             self.min_generation += 1
-            to_remove = [key for key in self.cache if self.used_generation[key] < self.min_generation]
+            active_keys = self._active_data_keys()
+            to_remove = [key for key in self.cache if key not in active_keys and self.used_generation[key] < self.min_generation]
             for key in to_remove:
                 del self.cache[key]
                 del self.used_generation[key]
@@ -546,8 +588,9 @@ class RAMPressureCache(LRUCache):
 
         clean_list = []
 
+        active_keys = self._active_data_keys()
         for key, cache_entry in self.cache.items():
-            if not free_active and self.used_generation[key] == self.generation:
+            if not free_active and (key in active_keys or self.used_generation[key] == self.generation):
                 continue
 
             if all_outputs_dynamic(cache_entry.outputs) and self.used_generation[key] == self.generation:

comfy_execution/graph_utils.py

@@ -1,3 +1,6 @@
+import contextvars
+
+
 def is_link(obj):
     if not isinstance(obj, list):
         return False
@@ -11,9 +14,7 @@ def is_link(obj):
 
 # The GraphBuilder is just a utility class that outputs graphs in the form expected by the ComfyUI back-end
 class GraphBuilder:
-    _default_prefix_root = ""
-    _default_prefix_call_index = 0
-    _default_prefix_graph_index = 0
+    _default_prefix = contextvars.ContextVar("graph_builder_default_prefix", default=("", 0, 0))
 
     def __init__(self, prefix = None):
         if prefix is None:
@@ -25,20 +26,19 @@ class GraphBuilder:
 
     @classmethod
     def set_default_prefix(cls, prefix_root, call_index, graph_index = 0):
-        cls._default_prefix_root = prefix_root
-        cls._default_prefix_call_index = call_index
-        cls._default_prefix_graph_index = graph_index
+        cls._default_prefix.set((prefix_root, call_index, graph_index))
 
     @classmethod
     def alloc_prefix(cls, root=None, call_index=None, graph_index=None):
+        default_root, default_call_index, default_graph_index = cls._default_prefix.get()
         if root is None:
-            root = GraphBuilder._default_prefix_root
+            root = default_root
         if call_index is None:
-            call_index = GraphBuilder._default_prefix_call_index
+            call_index = default_call_index
         if graph_index is None:
-            graph_index = GraphBuilder._default_prefix_graph_index
+            graph_index = default_graph_index
         result = f"{root}.{call_index}.{graph_index}."
-        GraphBuilder._default_prefix_graph_index += 1
+        cls._default_prefix.set((default_root, default_call_index, default_graph_index + 1))
         return result
 
     def node(self, class_type, id=None, **kwargs):

comfy_execution/progress.py

@@ -1,3 +1,4 @@
+import contextvars
 from typing import TypedDict, Dict, Optional, Tuple
 from typing_extensions import override
 from PIL import Image
@@ -11,6 +12,7 @@ from protocol import BinaryEventTypes
 from comfy_api import feature_flags
 
 PreviewImageTuple = Tuple[str, Image.Image, Optional[int]]
+_client_id_unset = object()
 
 class NodeState(Enum):
     Pending = "pending"
@@ -150,9 +152,15 @@ class WebUIProgressHandler(ProgressHandler):
     Handler that sends progress updates to the WebUI via WebSockets.
     """
 
-    def __init__(self, server_instance):
+    def __init__(self, server_instance, client_id=_client_id_unset):
         super().__init__("webui")
         self.server_instance = server_instance
+        self.client_id = client_id
+
+    def _client_id(self):
+        if self.client_id is _client_id_unset:
+            return self.server_instance.client_id
+        return self.client_id
 
     def set_registry(self, registry: "ProgressRegistry"):
         self.registry = registry
@@ -181,7 +189,7 @@ class WebUIProgressHandler(ProgressHandler):
         # Send a combined progress_state message with all node states
         # Include client_id to ensure message is only sent to the initiating client
         self.server_instance.send_sync(
-            "progress_state", {"prompt_id": prompt_id, "nodes": active_nodes}, self.server_instance.client_id
+            "progress_state", {"prompt_id": prompt_id, "nodes": active_nodes}, self._client_id()
         )
 
     @override
@@ -207,7 +215,7 @@ class WebUIProgressHandler(ProgressHandler):
             # Only send new format if client supports it
             if feature_flags.supports_feature(
                 self.server_instance.sockets_metadata,
-                self.server_instance.client_id,
+                self._client_id(),
                 "supports_preview_metadata",
             ):
                 metadata = {
@@ -224,7 +232,7 @@ class WebUIProgressHandler(ProgressHandler):
                 self.server_instance.send_sync(
                     BinaryEventTypes.PREVIEW_IMAGE_WITH_METADATA,
                     (image, metadata),
-                    self.server_instance.client_id,
+                    self._client_id(),
                 )
 
     @override
@@ -317,18 +325,19 @@ class ProgressRegistry:
         for handler in self.handlers.values():
             handler.reset()
 
-# Global registry instance
-global_progress_registry: ProgressRegistry | None = None
+progress_registry: contextvars.ContextVar[ProgressRegistry | None] = contextvars.ContextVar("progress_registry", default=None)
+
+def set_progress_registry(registry: ProgressRegistry):
+    return progress_registry.set(registry)
+
+def reset_progress_registry(token) -> None:
+    progress_registry.reset(token)
 
 def reset_progress_state(prompt_id: str, dynprompt: "DynamicPrompt") -> None:
-    global global_progress_registry
-
-    # Reset existing handlers if registry exists
-    if global_progress_registry is not None:
-        global_progress_registry.reset_handlers()
-
-    # Create new registry
-    global_progress_registry = ProgressRegistry(prompt_id, dynprompt)
+    current = progress_registry.get()
+    if current is not None:
+        current.reset_handlers()
+    progress_registry.set(ProgressRegistry(prompt_id, dynprompt))
 
 
 def add_progress_handler(handler: ProgressHandler) -> None:
@@ -338,11 +347,12 @@ def add_progress_handler(handler: ProgressHandler) -> None:
 
 
 def get_progress_state() -> ProgressRegistry:
-    global global_progress_registry
-    if global_progress_registry is None:
+    registry = progress_registry.get()
+    if registry is None:
         from comfy_execution.graph import DynamicPrompt
 
-        global_progress_registry = ProgressRegistry(
+        registry = ProgressRegistry(
             prompt_id="", dynprompt=DynamicPrompt({})
         )
-    return global_progress_registry
+        progress_registry.set(registry)
+    return registry

comfy_execution/utils.py

@@ -14,10 +14,25 @@ class ExecutionContext(NamedTuple):
     list_index: Optional[int]
 
 current_executing_context: contextvars.ContextVar[Optional[ExecutionContext]] = contextvars.ContextVar("current_executing_context", default=None)
+_client_id_unset = object()
+current_client_id = contextvars.ContextVar("current_client_id", default=_client_id_unset)
 
 def get_executing_context() -> Optional[ExecutionContext]:
     return current_executing_context.get(None)
 
+def get_current_client_id() -> Optional[str]:
+    value = current_client_id.get()
+    return None if value is _client_id_unset else value
+
+def has_current_client_id() -> bool:
+    return current_client_id.get() is not _client_id_unset
+
+def set_current_client_id(client_id: Optional[str]):
+    return current_client_id.set(client_id)
+
+def reset_current_client_id(token):
+    current_client_id.reset(token)
+
 class CurrentNodeContext:
     """
     Context manager for setting the current executing node context.

execution.py

@@ -10,6 +10,7 @@ import traceback
 from enum import Enum
 from typing import List, Literal, NamedTuple, Optional, Union
 import asyncio
+import contextlib
 
 import torch
 
@@ -40,7 +41,7 @@ from comfy_execution.graph import (
 from comfy_execution.graph_utils import GraphBuilder, is_link
 from comfy_execution.validation import validate_node_input
 from comfy_execution.progress import get_progress_state, reset_progress_state, add_progress_handler, WebUIProgressHandler
-from comfy_execution.utils import CurrentNodeContext
+from comfy_execution.utils import CurrentNodeContext, reset_current_client_id, set_current_client_id
 from comfy_execution.asset_enrichment import enrich_output_with_assets
 from comfy_api.internal import _ComfyNodeInternal, _NodeOutputInternal, first_real_override, is_class, make_locked_method_func
 from comfy_api.latest import io, _io
@@ -112,38 +113,41 @@ class CacheType(Enum):
 
 
 class CacheSet:
-    def __init__(self, cache_type=None, cache_args={}):
+    def __init__(self, cache_type=None, cache_args={}, outputs=None):
         if cache_type == CacheType.NONE:
-            self.init_null_cache()
-            logging.info("Disabling intermediate node cache.")
+            self.init_null_cache(outputs)
+            if outputs is None:
+                logging.info("Disabling intermediate node cache.")
         elif cache_type == CacheType.RAM_PRESSURE:
             cache_ram = cache_args.get("ram", 16.0)
-            self.init_ram_cache(cache_ram)
-            logging.info("Using RAM pressure cache.")
+            self.init_ram_cache(cache_ram, outputs)
+            if outputs is None:
+                logging.info("Using RAM pressure cache.")
         elif cache_type == CacheType.LRU:
             cache_size = cache_args.get("lru", 0)
-            self.init_lru_cache(cache_size)
-            logging.info("Using LRU cache")
+            self.init_lru_cache(cache_size, outputs)
+            if outputs is None:
+                logging.info("Using LRU cache")
         else:
-            self.init_classic_cache()
+            self.init_classic_cache(outputs)
 
         self.all = [self.outputs, self.objects]
 
     # Performs like the old cache -- dump data ASAP
-    def init_classic_cache(self):
-        self.outputs = HierarchicalCache(CacheKeySetInputSignature, enable_providers=True)
+    def init_classic_cache(self, outputs=None):
+        self.outputs = outputs if outputs is not None else HierarchicalCache(CacheKeySetInputSignature, enable_providers=True)
         self.objects = HierarchicalCache(CacheKeySetID)
 
-    def init_lru_cache(self, cache_size):
-        self.outputs = LRUCache(CacheKeySetInputSignature, max_size=cache_size, enable_providers=True)
+    def init_lru_cache(self, cache_size, outputs=None):
+        self.outputs = outputs if outputs is not None else LRUCache(CacheKeySetInputSignature, max_size=cache_size, enable_providers=True)
         self.objects = HierarchicalCache(CacheKeySetID)
 
-    def init_ram_cache(self, min_headroom):
-        self.outputs = RAMPressureCache(CacheKeySetInputSignature, enable_providers=True)
+    def init_ram_cache(self, min_headroom, outputs=None):
+        self.outputs = outputs if outputs is not None else RAMPressureCache(CacheKeySetInputSignature, enable_providers=True)
         self.objects = HierarchicalCache(CacheKeySetID)
 
-    def init_null_cache(self):
-        self.outputs = NullCache()
+    def init_null_cache(self, outputs=None):
+        self.outputs = outputs if outputs is not None else NullCache()
         self.objects = NullCache()
 
     def recursive_debug_dump(self):
@@ -425,15 +429,15 @@ def _is_intermediate_output(dynprompt, node_id):
     return getattr(class_def, 'HAS_INTERMEDIATE_OUTPUT', False)
 
 
-def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
+def _send_cached_ui(server, client_id, node_id, display_node_id, cached, prompt_id, ui_outputs):
     if cached.ui is not None:
         ui_outputs[node_id] = cached.ui
-    if server.client_id is None:
+    if client_id is None:
         return
     cached_ui = cached.ui or {}
-    server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
+    server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, client_id)
 
-async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
+async def execute(server, client_id, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
     unique_id = current_item
     real_node_id = dynprompt.get_real_node_id(unique_id)
     display_node_id = dynprompt.get_display_node_id(unique_id)
@@ -441,9 +445,12 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
     inputs = dynprompt.get_node(unique_id)['inputs']
     class_type = dynprompt.get_node(unique_id)['class_type']
     class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
+    prompt_queue = getattr(server, "prompt_queue", None)
+    if getattr(prompt_queue, "cooperative", False) and prompt_queue.is_cancelled(prompt_id):
+        return (ExecutionResult.FAILURE, {"node_id": real_node_id}, comfy.model_management.InterruptProcessingException())
     cached = await caches.outputs.get(unique_id)
     if cached is not None:
-        _send_cached_ui(server, unique_id, display_node_id, cached, prompt_id, ui_outputs)
+        _send_cached_ui(server, client_id, unique_id, display_node_id, cached, prompt_id, ui_outputs)
         get_progress_state().finish_progress(unique_id)
         execution_list.cache_update(unique_id, cached)
         return (ExecutionResult.SUCCESS, None, None)
@@ -489,9 +496,10 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
         else:
             get_progress_state().start_progress(unique_id)
             input_data_all, missing_keys, v3_data = get_input_data(inputs, class_def, unique_id, execution_list, dynprompt, extra_data)
-            if server.client_id is not None:
-                server.last_node_id = display_node_id
-                server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, server.client_id)
+            if client_id is not None:
+                if not getattr(getattr(server, "prompt_queue", None), "cooperative", False):
+                    server.last_node_id = display_node_id
+                server.send_sync("executing", { "node": unique_id, "display_node": display_node_id, "prompt_id": prompt_id }, client_id)
 
             obj = await caches.objects.get(unique_id)
             if obj is None:
@@ -531,12 +539,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                         "current_inputs": [],
                         "current_outputs": [],
                     }
-                    server.send_sync("execution_error", mes, server.client_id)
+                    server.send_sync("execution_error", mes, client_id)
                     return ExecutionBlocker(None)
                 else:
                     return block
             def pre_execute_cb(call_index):
-                # TODO - How to handle this with async functions without contextvars (which requires Python 3.12)?
                 GraphBuilder.set_default_prefix(unique_id, call_index, 0)
 
             try:
@@ -572,8 +579,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
                 },
                 "output": output_ui
             }
-            if server.client_id is not None:
-                server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)
+            if client_id is not None:
+                server.send_sync("executed", { "node": unique_id, "display_node": display_node_id, "output": output_ui, "prompt_id": prompt_id }, client_id)
         if has_subgraph:
             cached_outputs = []
             new_node_ids = []
@@ -660,16 +667,18 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
     return (ExecutionResult.SUCCESS, None, None)
 
 class PromptExecutor:
-    def __init__(self, server, cache_type=False, cache_args=None):
+    def __init__(self, server, cache_type=False, cache_args=None, shared_outputs=None):
         self.cache_args = cache_args
         self.cache_type = cache_type
         self.server = server
+        self.shared_outputs = shared_outputs
         self.reset()
 
     def reset(self):
-        self.caches = CacheSet(cache_type=self.cache_type, cache_args=self.cache_args)
+        self.caches = CacheSet(cache_type=self.cache_type, cache_args=self.cache_args, outputs=self.shared_outputs)
         self.status_messages = []
         self.success = True
+        self.client_id = None
 
     def add_message(self, event, data: dict, broadcast: bool):
         data = {
@@ -677,8 +686,8 @@ class PromptExecutor:
             "timestamp": int(time.time() * 1000),
         }
         self.status_messages.append((event, data))
-        if self.server.client_id is not None or broadcast:
-            self.server.send_sync(event, data, self.server.client_id)
+        if self.client_id is not None or broadcast:
+            self.server.send_sync(event, data, self.client_id)
 
     def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, error, ex):
         node_id = error["node_id"]
@@ -727,12 +736,14 @@ class PromptExecutor:
     async def execute_async(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
         set_preview_method(extra_data.get("preview_method"))
 
-        nodes.interrupt_processing(False)
+        cooperative = getattr(getattr(self.server, "prompt_queue", None), "cooperative", False)
+        if not cooperative:
+            nodes.interrupt_processing(False)
 
-        if "client_id" in extra_data:
-            self.server.client_id = extra_data["client_id"]
-        else:
-            self.server.client_id = None
+        self.client_id = extra_data.get("client_id")
+        if not cooperative:
+            self.server.client_id = self.client_id
+        client_id_token = set_current_client_id(self.client_id)
 
         self.status_messages = []
         self.add_message("execution_start", { "prompt_id": prompt_id}, broadcast=False)
@@ -741,13 +752,15 @@ class PromptExecutor:
         ram_headroom = int(self.cache_args["ram"] * (1024 ** 3))
         ram_inactive_headroom = int(self.cache_args["ram_inactive"] * (1024 ** 3))
         ram_release_callback = self.caches.outputs.ram_release if self.cache_type == CacheType.RAM_PRESSURE else None
-        comfy.memory_management.set_ram_cache_release_state(ram_release_callback, ram_headroom)
+        if not cooperative:
+            comfy.memory_management.set_ram_cache_release_state(ram_release_callback, ram_headroom)
 
         try:
-            with torch.inference_mode():
+            inference_context = contextlib.nullcontext() if cooperative else torch.inference_mode()
+            with inference_context:
                 dynamic_prompt = DynamicPrompt(prompt)
                 reset_progress_state(prompt_id, dynamic_prompt)
-                add_progress_handler(WebUIProgressHandler(self.server))
+                add_progress_handler(WebUIProgressHandler(self.server, self.client_id))
                 is_changed_cache = IsChangedCache(prompt_id, dynamic_prompt, self.caches.outputs)
                 for cache in self.caches.all:
                     await cache.set_prompt(dynamic_prompt, prompt.keys(), is_changed_cache)
@@ -782,7 +795,7 @@ class PromptExecutor:
                         break
 
                     assert node_id is not None, "Node ID should not be None at this point"
-                    result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
+                    result, error, ex = await execute(self.server, self.client_id, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
                     self.success = result != ExecutionResult.FAILURE
                     if result == ExecutionResult.FAILURE:
                         self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
@@ -792,7 +805,7 @@ class PromptExecutor:
                     else: # result == ExecutionResult.SUCCESS:
                         execution_list.complete_node_execution()
 
-                    if self.cache_type == CacheType.RAM_PRESSURE:
+                    if self.cache_type == CacheType.RAM_PRESSURE and not cooperative:
                         ram_release_callback(ram_inactive_headroom)
                         ram_shortfall = ram_headroom - psutil.virtual_memory().available
                         if ram_shortfall > 0:
@@ -816,7 +829,7 @@ class PromptExecutor:
                         cached = await self.caches.outputs.get(node_id)
                         if cached is not None:
                             display_node_id = dynamic_prompt.get_display_node_id(node_id)
-                            _send_cached_ui(self.server, node_id, display_node_id, cached, prompt_id, ui_node_outputs)
+                            _send_cached_ui(self.server, self.client_id, node_id, display_node_id, cached, prompt_id, ui_node_outputs)
                     self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
 
                 ui_outputs = {}
@@ -828,12 +841,17 @@ class PromptExecutor:
                     "outputs": ui_outputs,
                     "meta": meta_outputs,
                 }
-                self.server.last_node_id = None
+                if not cooperative:
+                    self.server.last_node_id = None
                 if comfy.model_management.DISABLE_SMART_MEMORY:
                     comfy.model_management.unload_all_models()
         finally:
-            comfy.memory_management.set_ram_cache_release_state(None, 0)
+            for cache in self.caches.all:
+                cache.release_prompt()
+            if not cooperative:
+                comfy.memory_management.set_ram_cache_release_state(None, 0)
             self._notify_prompt_lifecycle("end", prompt_id)
+            reset_current_client_id(client_id_token)
 
 
 async def validate_inputs(prompt_id, prompt, item, validated, visiting=None):
@@ -1249,6 +1267,9 @@ class PromptQueue:
         self.task_counter = 0
         self.queue = []
         self.currently_running = {}
+        self.cancelled_prompts = set()
+        self.cooperative = False
+        self.cooperative_drain = False
         self.history = {}
         self.flags = {}
 
@@ -1271,6 +1292,17 @@ class PromptQueue:
             self.server.queue_updated()
             return (item, i)
 
+    def get_if(self, predicate):
+        with self.mutex:
+            if self.cooperative_drain or len(self.queue) == 0 or not predicate(self.queue[0]):
+                return None
+            item = heapq.heappop(self.queue)
+            i = self.task_counter
+            self.currently_running[i] = copy.deepcopy(item)
+            self.task_counter += 1
+            self.server.queue_updated()
+            return (item, i)
+
     class ExecutionStatus(NamedTuple):
         status_str: Literal['success', 'error']
         completed: bool
@@ -1280,6 +1312,7 @@ class PromptQueue:
                   status: Optional['PromptQueue.ExecutionStatus'], process_item=None):
         with self.mutex:
             prompt = self.currently_running.pop(item_id)
+            self.cancelled_prompts.discard(prompt[1])
             if len(self.history) > MAXIMUM_HISTORY_SIZE:
                 self.history.pop(next(iter(self.history)))
 
@@ -1316,22 +1349,51 @@ class PromptQueue:
     def interrupt_if_running(self, prompt_id):
         """Interrupt the running prompt with this id, atomically.
 
-        Checks the live running set and signals the interrupt under the queue
-        mutex, so the worker cannot move the job to done (and start the next
-        prompt) in between. Returns True if a matching job was running and an
-        interrupt was signalled, False otherwise. The atomicity is what keeps a
-        cancel from landing on an unrelated prompt that started after a separate
-        is-running check: the global interrupt flag is reset at the start of
-        every prompt (execute_async), so a job that finishes before consuming
-        the flag cannot leak the interrupt onto its successor.
+        Cooperative workers use prompt-scoped cancellation when other prompts
+        are active. Serial execution and single-prompt cooperative execution can
+        safely use the global interrupt flag as well. Cancellation is best-effort:
+        if execution has already completed before observing the marker, task_done
+        records the completed result and clears the marker.
         """
         with self.mutex:
             for item in self.currently_running.values():
                 if item[1] == prompt_id:
-                    nodes.interrupt_processing()
+                    if self.cooperative:
+                        self.cancelled_prompts.add(prompt_id)
+                        if len(self.currently_running) == 1:
+                            self.cooperative_drain = True
+                            nodes.interrupt_processing()
+                    else:
+                        nodes.interrupt_processing()
                     return True
         return False
 
+    def interrupt_all_running(self):
+        with self.mutex:
+            if self.cooperative and self.currently_running:
+                self.cancelled_prompts.update(item[1] for item in self.currently_running.values())
+                self.cooperative_drain = True
+            nodes.interrupt_processing()
+
+    def finish_cooperative_drain(self):
+        with self.mutex:
+            self.cooperative_drain = False
+
+    def is_cooperative_draining(self):
+        with self.mutex:
+            return self.cooperative_drain
+
+    def set_cooperative(self, enabled):
+        with self.mutex:
+            self.cooperative = enabled
+            if not enabled:
+                self.cooperative_drain = False
+                self.cancelled_prompts.clear()
+
+    def is_cancelled(self, prompt_id):
+        with self.mutex:
+            return prompt_id in self.cancelled_prompts
+
     def get_tasks_remaining(self):
         with self.mutex:
             return len(self.queue) + len(self.currently_running)

main.py

@@ -29,10 +29,13 @@ import logging
 import signal
 import sys
 from comfy_execution.progress import get_progress_state
-from comfy_execution.utils import get_executing_context
+from comfy_execution.utils import get_current_client_id, get_executing_context, has_current_client_id
 from comfy_api import feature_flags
 from app.database.db import init_db, dependencies_available
 
+if args.continuous_batching:
+    logging.info("Continuous batching enabled; DynamicVRAM is disabled and the legacy ModelPatcher will be used")
+
 if __name__ == "__main__":
     #NOTE: These do not do anything on core ComfyUI, they are for custom nodes.
     os.environ['HF_HUB_DISABLE_TELEMETRY'] = '1'
@@ -225,6 +228,7 @@ import gc
 if 'torch' in sys.modules:
     logging.warning("WARNING: Potential Error in code: Torch already imported, torch should never be imported before this point.")
 
+import torch
 
 import comfy.utils
 
@@ -232,6 +236,8 @@ import execution
 import server
 from protocol import BinaryEventTypes
 import nodes
+import comfy.continuous_batching
+import comfy.text_encoders.anima_cache
 import comfy.model_management
 import comfyui_version
 import app.logger
@@ -313,8 +319,7 @@ def _collect_output_absolute_paths(history_result: dict) -> list[str]:
     return paths
 
 
-def prompt_worker(q, server_instance):
-    current_time: float = 0.0
+def prompt_executor_config():
     cache_ram = 0
     cache_ram_inactive = 0
     if not args.cache_classic and not args.cache_none and args.cache_lru <= 0:
@@ -333,7 +338,189 @@ def prompt_worker(q, server_instance):
     elif args.cache_none:
         cache_type = execution.CacheType.NONE
 
-    e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args={ "lru" : args.cache_lru, "ram" : cache_ram, "ram_inactive" : cache_ram_inactive } )
+    return cache_type, {"lru": args.cache_lru, "ram": cache_ram, "ram_inactive": cache_ram_inactive}
+
+
+async def execute_prompt_async(q, server_instance, item, item_id, cache_type, cache_args, shared_outputs):
+    executor = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args=cache_args, shared_outputs=shared_outputs)
+    execution_start_time = time.perf_counter()
+    prompt_id = item[1]
+    server_instance.last_prompt_id = prompt_id
+
+    sensitive = item[5]
+    extra_data = item[3].copy()
+    for key in sensitive:
+        extra_data[key] = sensitive[key]
+
+    asset_seeder.pause()
+    await executor.execute_async(item[2], prompt_id, extra_data, item[4])
+
+    remove_sensitive = lambda prompt: prompt[:5] + prompt[6:]
+    q.task_done(
+        item_id,
+        executor.history_result,
+        status=execution.PromptQueue.ExecutionStatus(
+            status_str="success" if executor.success else "error",
+            completed=executor.success,
+            messages=executor.status_messages,
+        ),
+        process_item=remove_sensitive,
+    )
+    client_id = extra_data.get("client_id")
+    if client_id is not None:
+        server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, client_id)
+
+    execution_time = time.perf_counter() - execution_start_time
+    if execution_time > 600:
+        execution_time = time.strftime("%H:%M:%S", time.gmtime(execution_time))
+        logging.info(f"Prompt executed in {execution_time}", extra={"color": "green"})
+    else:
+        logging.info("Prompt executed in {:.2f} seconds".format(execution_time), extra={"color": "green"})
+
+    if not asset_seeder.is_disabled():
+        paths = _collect_output_absolute_paths(executor.history_result)
+        register_output_files(paths, job_id=prompt_id)
+
+
+def _freeze_prompt_value(value):
+    if isinstance(value, dict):
+        return tuple((key, _freeze_prompt_value(item)) for key, item in sorted(value.items()))
+    if isinstance(value, list):
+        return tuple(_freeze_prompt_value(item) for item in value)
+    return value
+
+
+def _prompt_dependency_signature(prompt, node_id, memo):
+    if node_id in memo:
+        return memo[node_id]
+    node = prompt[node_id]
+    inputs = []
+    for name, value in sorted(node.get("inputs", {}).items()):
+        if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in prompt and isinstance(value[1], (int, float)):
+            value = ("link", value[1], _prompt_dependency_signature(prompt, value[0], memo))
+        else:
+            value = ("value", _freeze_prompt_value(value))
+        inputs.append((name, value))
+    signature = (node["class_type"], tuple(inputs))
+    memo[node_id] = signature
+    return signature
+
+
+def continuous_prompt_key(item):
+    prompt = item[2]
+    outputs_to_execute = item[4]
+    pending = list(outputs_to_execute) if isinstance(outputs_to_execute, (list, tuple, set)) else [outputs_to_execute]
+    dependencies = set()
+    while pending:
+        node_id = pending.pop()
+        if node_id in dependencies:
+            continue
+        node = prompt.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        dependencies.add(node_id)
+        for value in node.get("inputs", {}).values():
+            if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str) and value[0] in prompt and isinstance(value[1], (int, float)):
+                pending.append(value[0])
+
+    sampler_ids = [node_id for node_id in dependencies if prompt[node_id].get("class_type") in comfy.continuous_batching.CONTINUOUS_SAMPLER_NODE_FAMILIES]
+    if len(sampler_ids) != 1:
+        return None
+    sampler_type = prompt[sampler_ids[0]]["class_type"]
+    model_input = prompt[sampler_ids[0]].get("inputs", {}).get("model")
+    if not isinstance(model_input, list) or len(model_input) != 2 or model_input[0] not in prompt:
+        return None
+    return (
+        sampler_type,
+        _prompt_dependency_signature(prompt, model_input[0], {}),
+        item[3].get("preview_method"),
+    )
+
+
+async def cooperative_prompt_worker(q, server_instance, max_prompts):
+    cache_type, cache_args = prompt_executor_config()
+    shared_outputs = execution.CacheSet(cache_type=cache_type, cache_args=cache_args).outputs
+    active = set()
+    active_key = None
+    worker_cache_scope = None
+    group_cache_scope = None
+    q.set_cooperative(True)
+    comfy.continuous_batching.set_cancel_checker(q.is_cancelled)
+    ram_headroom = int(cache_args["ram"] * (1024 ** 3))
+    ram_inactive_headroom = int(cache_args["ram_inactive"] * (1024 ** 3))
+    ram_release_callback = shared_outputs.ram_release if cache_type == execution.CacheType.RAM_PRESSURE else None
+    comfy.memory_management.set_ram_cache_release_state(ram_release_callback, ram_headroom)
+    try:
+        worker_cache_scope = comfy.text_encoders.anima_cache.begin_cache_scope(False)
+        # PromptExecutor originally owned this scope. It is moved to one worker owner because
+        # thread-local inference mode can break when interleaved async tasks exit out of order.
+        with torch.inference_mode():
+            while True:
+                while len(active) < max_prompts:
+                    queue_item = q.get_if(lambda item: not active or active_key is not None and continuous_prompt_key(item) == active_key)
+                    if queue_item is None:
+                        break
+                    item, item_id = queue_item
+                    item_key = continuous_prompt_key(item)
+                    if not active:
+                        nodes.interrupt_processing(False)
+                        active_key = item_key
+                        family = comfy.continuous_batching.CONTINUOUS_SAMPLER_NODE_FAMILIES.get(item_key[0]) if item_key is not None else None
+                        group_cache_scope = comfy.text_encoders.anima_cache.begin_cache_scope(family == comfy.continuous_batching.FAMILY_ANIMA)
+                    active.add(asyncio.create_task(execute_prompt_async(q, server_instance, item, item_id, cache_type, cache_args, shared_outputs)))
+                    if item_key is None:
+                        break
+
+                if active:
+                    done, active = await asyncio.wait(active, timeout=0.05, return_when=asyncio.FIRST_COMPLETED)
+                    for task in done:
+                        task.result()
+                    if done and not active:
+                        active_key = None
+                        comfy.text_encoders.anima_cache.end_cache_scope(group_cache_scope)
+                        group_cache_scope = None
+                        q.finish_cooperative_drain()
+                        if ram_release_callback is not None:
+                            ram_release_callback(ram_inactive_headroom, free_active=True)
+                        gc.collect()
+                        comfy.model_management.soft_empty_cache()
+                        hook_breaker_ac10a0.restore_functions()
+                        if not asset_seeder.is_disabled():
+                            asset_seeder.enqueue_enrich(roots=("output",), compute_hashes=args.enable_asset_hashing)
+                        asset_seeder.resume()
+                else:
+                    await asyncio.sleep(0.05)
+
+                if not active:
+                    flags = q.get_flags()
+                    free_memory = flags.get("free_memory", False)
+                    if flags.get("unload_models", free_memory):
+                        comfy.model_management.unload_all_models()
+                    if free_memory:
+                        gc.collect()
+                        comfy.model_management.soft_empty_cache()
+    finally:
+        try:
+            if group_cache_scope is not None:
+                comfy.text_encoders.anima_cache.end_cache_scope(group_cache_scope)
+        finally:
+            if worker_cache_scope is not None:
+                comfy.text_encoders.anima_cache.end_cache_scope(worker_cache_scope)
+        comfy.memory_management.set_ram_cache_release_state(None, 0)
+        comfy.continuous_batching.set_cancel_checker(None)
+        q.set_cooperative(False)
+
+
+def prompt_worker(q, server_instance):
+    max_prompts = args.continuous_batching
+    if max_prompts:
+        asyncio.run(cooperative_prompt_worker(q, server_instance, max_prompts))
+        return
+
+    current_time: float = 0.0
+    cache_type, cache_args = prompt_executor_config()
+
+    e = execution.PromptExecutor(server_instance, cache_type=cache_type, cache_args=cache_args)
     last_gc_collect = 0
     need_gc = False
     gc_collect_interval = 10.0
@@ -367,8 +554,8 @@ def prompt_worker(q, server_instance):
                             status_str='success' if e.success else 'error',
                             completed=e.success,
                             messages=e.status_messages), process_item=remove_sensitive)
-            if server_instance.client_id is not None:
-                server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
+            if e.client_id is not None:
+                server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, e.client_id)
 
             current_time = time.perf_counter()
             execution_time = current_time - execution_start_time
@@ -434,18 +621,21 @@ def hijack_progress(server_instance):
         progress = {"value": value, "max": total, "prompt_id": prompt_id, "node": node_id}
         get_progress_state().update_progress(node_id, value, total, preview_image)
 
-        server_instance.send_sync("progress", progress, server_instance.client_id)
+        client_id = get_current_client_id()
+        if not has_current_client_id():
+            client_id = server_instance.client_id
+        server_instance.send_sync("progress", progress, client_id)
         if preview_image is not None:
             # Only send old method if client doesn't support preview metadata
             if not feature_flags.supports_feature(
                 server_instance.sockets_metadata,
-                server_instance.client_id,
+                client_id,
                 "supports_preview_metadata",
             ):
                 server_instance.send_sync(
                     BinaryEventTypes.UNENCODED_PREVIEW_IMAGE,
                     preview_image,
-                    server_instance.client_id,
+                    client_id,
                 )
 
     comfy.utils.set_progress_bar_global_hook(hook)

nodes.py

@@ -26,7 +26,10 @@ import comfy.sample
 import comfy.sd
 import comfy.utils
 import comfy.controlnet
+import comfy.continuous_batching
 from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict, FileLocator
+from comfy_execution.progress import get_progress_state
+from comfy_execution.utils import get_current_client_id, get_executing_context
 from comfy_api.internal import register_versions, ComfyAPIWithVersion
 from comfy_api.version_list import supported_versions
 from comfy_api.latest import io, ComfyExtension, InputImpl
@@ -1606,6 +1609,87 @@ class KSampler:
     def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
         return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
 
+
+class _ContinuousKSampler:
+    FAMILY = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
+                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
+                "positive": ("CONDITIONING",),
+                "negative": ("CONDITIONING",),
+                "latent_image": ("LATENT",),
+                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "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 = ("LATENT",)
+    OUTPUT_TOOLTIPS = ("The denoised latent.",)
+    FUNCTION = "sample"
+    CATEGORY = "model/sampling"
+
+    async def sample(self, model, seed, steps, cfg, scheduler, positive, negative, latent_image, denoise=1.0, max_batch_size=4, admission_delay_ms=2.0):
+        latent = latent_image["samples"]
+        latent = comfy.sample.fix_empty_latent_channels(model, latent, latent_image.get("downscale_ratio_spacial"), latent_image.get("downscale_ratio_temporal"))
+        sampler = comfy.samplers.KSampler(model, steps=steps, device=model.load_device, sampler="euler", scheduler=scheduler, denoise=denoise, model_options=model.model_options)
+        if len(sampler.sigmas) == 0:
+            out = latent_image.copy()
+            out.pop("downscale_ratio_spacial", None)
+            out.pop("downscale_ratio_temporal", None)
+            out["samples"] = latent
+            return (out,)
+        if "noise_mask" in latent_image:
+            raise ValueError("Continuous batching does not support noise masks")
+        noise = comfy.sample.prepare_noise(latent, seed, latent_image.get("batch_index"))
+        execution_context = get_executing_context()
+        state = comfy.continuous_batching.ContinuousBatchRequest(
+            family=self.FAMILY,
+            model_patcher=model,
+            noise=noise,
+            latent_image=latent,
+            positive=positive,
+            negative=negative,
+            sigmas=sampler.sigmas,
+            callback=latent_preview.prepare_callback(model, steps),
+            seed=seed,
+            cfg=cfg,
+            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,
+            node_id=execution_context.node_id if execution_context is not None else None,
+            client_id=get_current_client_id(),
+            progress_registry=get_progress_state(),
+        )
+        samples = await comfy.continuous_batching.sample_euler_continuous(state)
+        out = latent_image.copy()
+        out.pop("downscale_ratio_spacial", None)
+        out.pop("downscale_ratio_temporal", None)
+        out["samples"] = samples
+        return (out,)
+
+
+class AnimaContinuousKSampler(_ContinuousKSampler):
+    FAMILY = comfy.continuous_batching.FAMILY_ANIMA
+    DESCRIPTION = "Continuously batches compatible Anima requests at Euler denoising-step boundaries."
+
+
+class SD15ContinuousKSampler(_ContinuousKSampler):
+    FAMILY = comfy.continuous_batching.FAMILY_SD15
+    DESCRIPTION = "Continuously batches compatible SD1.5 requests at Euler denoising-step boundaries."
+
+
+class SDXLContinuousKSampler(_ContinuousKSampler):
+    FAMILY = comfy.continuous_batching.FAMILY_SDXL
+    DESCRIPTION = "Continuously batches compatible SDXL requests at Euler denoising-step boundaries."
+
 class KSamplerAdvanced:
     @classmethod
     def INPUT_TYPES(s):
@@ -2048,6 +2132,9 @@ class ImagePadForOutpaint:
 
 NODE_CLASS_MAPPINGS = {
     "KSampler": KSampler,
+    "AnimaContinuousKSampler": AnimaContinuousKSampler,
+    "SD15ContinuousKSampler": SD15ContinuousKSampler,
+    "SDXLContinuousKSampler": SDXLContinuousKSampler,
     "CheckpointLoaderSimple": CheckpointLoaderSimple,
     "CLIPTextEncode": CLIPTextEncode,
     "CLIPSetLastLayer": CLIPSetLastLayer,
@@ -2120,6 +2207,9 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     # Sampling
     "KSampler": "KSampler",
+    "AnimaContinuousKSampler": "Anima Continuous KSampler",
+    "SD15ContinuousKSampler": "SD1.5 Continuous KSampler",
+    "SDXLContinuousKSampler": "SDXL Continuous KSampler",
     "KSamplerAdvanced": "KSampler (Advanced)",
     # Loaders
     "CheckpointLoader": "Load Checkpoint With Config (DEPRECATED)",

server.py

@@ -1157,25 +1157,14 @@ class PromptServer():
             # Check if a specific prompt_id was provided for targeted interruption
             prompt_id = json_data.get('prompt_id')
             if prompt_id:
-                currently_running, _ = self.prompt_queue.get_current_queue()
-
-                # Check if the prompt_id matches any currently running prompt
-                should_interrupt = False
-                for item in currently_running:
-                    # item structure: (number, prompt_id, prompt, extra_data, outputs_to_execute)
-                    if item[1] == prompt_id:
-                        logging.info(f"Interrupting prompt {prompt_id}")
-                        should_interrupt = True
-                        break
-
-                if should_interrupt:
-                    nodes.interrupt_processing()
+                if self.prompt_queue.interrupt_if_running(prompt_id):
+                    logging.info(f"Interrupting prompt {prompt_id}")
                 else:
                     logging.info(f"Prompt {prompt_id} is not currently running, skipping interrupt")
             else:
                 # No prompt_id provided, do a global interrupt
                 logging.info("Global interrupt (no prompt_id specified)")
-                nodes.interrupt_processing()
+                self.prompt_queue.interrupt_all_running()
 
             return web.Response(status=200)
 

tests-unit/comfy_test/test_anima_cache.py

@@ -0,0 +1,261 @@
+import asyncio
+
+import pytest
+import torch
+
+from comfy.text_encoders import anima_cache
+
+
+def tokens(*ids):
+    return [list(ids)]
+
+
+class TinyCausalTransformer:
+    def __init__(self):
+        self.calls = []
+        self.num_tokens = []
+
+    def __call__(self, _, attention_mask, embeds, num_tokens, intermediate_output, final_layer_norm_intermediate, dtype, embeds_info, past_key_values=None):
+        self.calls.append((embeds.shape[1], past_key_values is not None))
+        self.num_tokens.append(tuple(num_tokens))
+        if past_key_values and embeds.shape[1] > 1:
+            raise RuntimeError("cached multi-token suffix would use an invalid causal mask")
+        prefix = 0
+        if past_key_values:
+            prefix = past_key_values[0][2]
+            previous = past_key_values[0][0][:, :, :prefix].reshape(embeds.shape[0], prefix, 1)
+        else:
+            previous = embeds[:, :0]
+        sequence = torch.cat((previous, embeds), dim=1)
+        hidden = sequence.cumsum(dim=1)[:, prefix:]
+        key = sequence.reshape(sequence.shape[0], 1, sequence.shape[1], 1).clone()
+        value = (sequence * 2).reshape(sequence.shape[0], 1, sequence.shape[1], 1).clone()
+        if past_key_values is None:
+            return hidden, None
+        return hidden, None, [(key, value, sequence.shape[1])]
+
+
+class CacheOwner:
+    def __init__(self, weight_uuid="weights-a"):
+        self.current_weight_patches_uuid = weight_uuid
+
+
+def cached_forward(transformer, token_ids, dtype=torch.float32, attention_mask=None, intermediate_output=None, embeds_info=None, owner=None):
+    embeds = torch.tensor(token_ids, dtype=dtype).reshape(1, -1, 1)
+    owner = owner or CacheOwner()
+    return anima_cache.forward(transformer, owner, tokens(*token_ids), attention_mask, embeds, [len(token_ids)], intermediate_output, False, torch.float32, embeds_info or [])
+
+
+@pytest.fixture
+def prefix_cache():
+    scope = anima_cache.begin_cache_scope()
+    try:
+        yield scope[0]
+    finally:
+        anima_cache.end_cache_scope(scope)
+
+
+def test_token_ids_accept_runtime_lists_tensors_and_unit_weight_legacy_pairs():
+    assert anima_cache._token_ids([[1, 2, 3]]) == (1, 2, 3)
+    assert anima_cache._token_ids(torch.tensor([[1, 2, 3]])) == (1, 2, 3)
+    assert anima_cache._token_ids([[(1, 1.0), (2, 1)]]) == (1, 2)
+    assert anima_cache._token_ids([[(1, 0.5), (2, 1.0)]]) is None
+    assert anima_cache._token_ids([[(1, 1.0, "custom")]]) is None
+    assert anima_cache._token_ids([[{"type": "embedding"}]]) is None
+
+
+def test_non_unit_legacy_weight_bypasses_cache(prefix_cache):
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    embeds = torch.tensor([1.0, 2.0]).reshape(1, -1, 1)
+    weighted_tokens = [[(1, 0.5), (2, 1.0)]]
+
+    anima_cache.forward(transformer, owner, weighted_tokens, None, embeds, [2], None, False, torch.float32, [])
+    anima_cache.forward(transformer, owner, weighted_tokens, None, embeds, [2], None, False, torch.float32, [])
+
+    assert transformer.calls == [(2, False), (2, False)]
+    assert len(prefix_cache) == 0
+
+
+def test_reuses_prefix_for_extension_exactly(caplog, prefix_cache):
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    cached_forward(transformer, [1, 2, 3], owner=owner)
+    with caplog.at_level("DEBUG"):
+        output, _ = cached_forward(transformer, [1, 2, 3, 4], owner=owner)
+
+    assert torch.equal(output, torch.tensor([[[1.0], [3.0], [6.0], [10.0]]]))
+    assert transformer.calls == [(3, True), (1, True)]
+    assert transformer.num_tokens == [(3,), (1,)]
+    assert "Anima Qwen cache reused 3 prefix tokens" in caplog.messages
+
+
+def test_exact_hit_does_not_call_transformer(prefix_cache):
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    expected, _ = cached_forward(transformer, [1, 2, 3], owner=owner)
+    call_count = len(transformer.calls)
+
+    output, _ = cached_forward(transformer, [1, 2, 3], owner=owner)
+
+    assert torch.equal(output, expected)
+    assert len(transformer.calls) == call_count
+
+
+def test_strict_prefix_of_cached_prompt_is_copied_without_transformer_call(prefix_cache):
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    cached_forward(transformer, [1, 2, 3, 4], owner=owner)
+    output, _ = cached_forward(transformer, [1, 2], owner=owner)
+
+    assert torch.equal(output, torch.tensor([[[1.0], [3.0]]]))
+    assert transformer.calls == [(4, True)]
+    cached = prefix_cache[owner]
+    assert cached[2] == (1, 2)
+    assert cached[3]._base is None
+    assert all(key.shape[2] == value.shape[2] == index == 2 for key, value, index in cached[4])
+    assert all(key._base is None and value._base is None for key, value, _ in cached[4])
+
+
+def test_diverging_suffix_matches_full_causal_forward(prefix_cache):
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    cached_forward(transformer, [1, 2, 9], owner=owner)
+    output, _ = cached_forward(transformer, [1, 2, 4, 5], owner=owner)
+
+    assert torch.equal(output, torch.tensor([[[1.0], [3.0], [7.0], [12.0]]]))
+    assert transformer.calls == [(3, True), (1, True), (1, True)]
+    assert transformer.num_tokens == [(3,), (1,), (1,)]
+
+
+def test_cache_is_isolated_by_owner_and_transformer_identity(prefix_cache):
+    shared_transformer = TinyCausalTransformer()
+    first_owner = CacheOwner("same-uuid")
+    second_owner = CacheOwner("same-uuid")
+    cached_forward(shared_transformer, [1, 2], owner=first_owner)
+    cached_forward(shared_transformer, [1, 2, 3], owner=second_owner)
+    assert shared_transformer.calls == [(2, True), (3, True)]
+
+    replacement = TinyCausalTransformer()
+    cached_forward(replacement, [1, 2, 3], owner=first_owner)
+    assert replacement.calls == [(3, True)]
+
+
+def test_scope_end_clears_cached_prefixes_and_disables_reuse():
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    scope = anima_cache.begin_cache_scope()
+    cache = scope[0]
+    try:
+        cached_forward(transformer, [1, 2], owner=owner)
+        assert len(cache) == 1
+    finally:
+        anima_cache.end_cache_scope(scope)
+
+    assert len(cache) == 0
+    cached_forward(transformer, [1, 2, 3], owner=owner)
+    assert transformer.calls[-1] == (3, False)
+    assert len(cache) == 0
+
+
+def test_child_tasks_share_the_same_cache_scope():
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+
+    async def run():
+        scope = anima_cache.begin_cache_scope()
+        primed = asyncio.Event()
+
+        async def prime():
+            cached_forward(transformer, [1, 2], owner=owner)
+            primed.set()
+
+        async def reuse():
+            await primed.wait()
+            return cached_forward(transformer, [1, 2, 3], owner=owner)
+
+        try:
+            first = asyncio.create_task(prime())
+            second = asyncio.create_task(reuse())
+            await first
+            return await second
+        finally:
+            anima_cache.end_cache_scope(scope)
+
+    output, _ = asyncio.run(run())
+
+    assert torch.equal(output, torch.tensor([[[1.0], [3.0], [6.0]]]))
+    assert transformer.calls == [(2, True), (1, True)]
+
+
+def test_nested_cache_scopes_do_not_share_prefixes():
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    outer_scope = anima_cache.begin_cache_scope()
+    outer_cache = outer_scope[0]
+    try:
+        cached_forward(transformer, [1, 2], owner=owner)
+        inner_scope = anima_cache.begin_cache_scope()
+        inner_cache = inner_scope[0]
+        try:
+            cached_forward(transformer, [1, 2, 3], owner=owner)
+            assert len(inner_cache) == 1
+        finally:
+            anima_cache.end_cache_scope(inner_scope)
+
+        assert len(inner_cache) == 0
+        cached_forward(transformer, [1, 2, 4], owner=owner)
+        assert len(outer_cache) == 1
+    finally:
+        anima_cache.end_cache_scope(outer_scope)
+
+    assert transformer.calls == [(2, True), (3, True), (1, True)]
+
+
+def test_disabled_scope_hides_an_outer_cache_scope():
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    outer_scope = anima_cache.begin_cache_scope()
+    try:
+        cached_forward(transformer, [1, 2], owner=owner)
+        disabled_scope = anima_cache.begin_cache_scope(False)
+        try:
+            cached_forward(transformer, [1, 2], owner=owner)
+        finally:
+            anima_cache.end_cache_scope(disabled_scope)
+
+        call_count = len(transformer.calls)
+        cached_forward(transformer, [1, 2], owner=owner)
+        assert len(transformer.calls) == call_count
+    finally:
+        anima_cache.end_cache_scope(outer_scope)
+
+    assert transformer.calls == [(2, True), (2, False)]
+
+
+def test_unsafe_inputs_and_changed_model_state_do_not_reuse_cache(prefix_cache):
+    transformer = TinyCausalTransformer()
+    owner = CacheOwner()
+    cached_forward(transformer, [1, 2], attention_mask=torch.tensor([[True, False]]), owner=owner)
+    cached_forward(transformer, [1, 2], intermediate_output=0, owner=owner)
+    cached_forward(transformer, [1, 2], embeds_info=[{"custom": True}], owner=owner)
+    assert transformer.calls == [(2, False), (2, False), (2, False)]
+    assert len(prefix_cache) == 0
+
+    cached_forward(transformer, [1, 2], dtype=torch.float32, owner=owner)
+    cached_forward(transformer, [1, 2, 3], dtype=torch.float64, owner=owner)
+    owner.current_weight_patches_uuid = "weights-b"
+    cached_forward(transformer, [1, 2, 3, 4], dtype=torch.float64, owner=owner)
+    assert transformer.calls[-3:] == [(2, True), (3, True), (4, True)]
+
+
+def test_cache_owner_reference_is_not_registered_as_child_module():
+    owner = torch.nn.Module()
+    child = torch.nn.Module()
+    owner.add_module("child", child)
+
+    anima_cache.set_owner(child, owner)
+
+    assert anima_cache.get_owner(child) is owner
+    assert "_anima_cache_owner" not in child._modules
+    assert list(owner.state_dict()) == []

tests-unit/comfy_test/test_continuous_batching.py

@@ -0,0 +1,425 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import comfy.conds
+import comfy.model_base
+import comfy.patcher_extension
+import comfy.supported_models
+from comfy.continuous_batching import (
+    FAMILY_ANIMA,
+    FAMILY_SD15,
+    FAMILY_SDXL,
+    ContinuousBatchCoordinator,
+    ContinuousBatchSession,
+    _cfg_branches,
+    _conditioning_structure,
+    _processed_conditioning_signature,
+    _validate_conditioning,
+    _validate_model_extensions,
+    _validate_model_family,
+    cfg_combine,
+    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
+
+
+class FakeState:
+    def __init__(self, name, steps, max_batch_size=2, admission_delay=0.0):
+        self.name = name
+        self.steps = steps
+        self.max_batch_size = max_batch_size
+        self.admission_delay = admission_delay
+        self.model_patcher = object()
+        self.family = "test"
+        self.output = None
+        self.cleared = False
+
+    def clear(self):
+        self.cleared = True
+
+    def key(self):
+        return "group"
+
+
+class FakeSession:
+    def __init__(self, fail=False, on_first_step=None):
+        self.batches = []
+        self.closed = 0
+        self.open = False
+        self.fail = fail
+        self.on_first_step = on_first_step
+        self.reload_requests = 0
+
+    def step(self, states):
+        self.open = True
+        self.batches.append([state.name for state in states])
+        if len(self.batches) == 1 and self.on_first_step is not None:
+            self.on_first_step()
+        if self.fail:
+            raise RuntimeError("denoiser failed")
+        updates = []
+        for state in states:
+            state.steps -= 1
+            finished = state.steps == 0
+            if finished:
+                state.output = state.name + " output"
+            updates.append((state, finished))
+        return updates
+
+    def close(self):
+        self.closed += 1
+        self.open = False
+
+    def request_model_reload(self):
+        self.reload_requests += 1
+
+
+def _model_patcher(model):
+    return SimpleNamespace(model=model)
+
+
+def _bare_model(model_type, config_type, in_channels=4, latent_channels=4, concat_keys=()):
+    model = object.__new__(model_type)
+    torch.nn.Module.__init__(model)
+    model.model_config = object.__new__(config_type)
+    model.diffusion_model = SimpleNamespace(in_channels=in_channels)
+    model.latent_format = SimpleNamespace(latent_channels=latent_channels)
+    model.concat_keys = concat_keys
+    return model
+
+
+def _extension_patcher(model_options=None, callbacks=None, additional_models=None, wrappers=None):
+    return SimpleNamespace(
+        model_options=model_options or {},
+        callbacks=callbacks or {},
+        additional_models=additional_models or {},
+        wrappers=wrappers or {},
+    )
+
+
+def _processed_cond(conditioning):
+    return SimpleNamespace(area=None, control=None, patches=None, hooks=None, conditioning=conditioning)
+
+
+def _batch_cond(x, length, marker, uuid):
+    cond = _processed_cond({"c_crossattn": comfy.conds.CONDCrossAttn(torch.full((1, length, 4), marker))})
+    cond.input_x = x
+    cond.uuid = uuid
+    return cond
+
+
+def _batch_state(sigma, cfg, negative_marker, positive_marker):
+    x = torch.zeros(1, 4, 2, 2)
+    return SimpleNamespace(
+        family=FAMILY_SD15,
+        x=x,
+        sigmas=torch.tensor([sigma, 0.0]),
+        index=0,
+        cfg=cfg,
+        conds={
+            "negative": [_batch_cond(x, 77, negative_marker, f"negative-{negative_marker}")],
+            "positive": [_batch_cond(x, 154, positive_marker, f"positive-{positive_marker}")],
+        },
+    )
+
+
+class _RecordingPatcher:
+    def __init__(self):
+        self.prepared_sigmas = []
+
+    def prepare_state(self, sigmas, model_options):
+        self.prepared_sigmas.append(sigmas.clone())
+
+    def apply_hooks(self, hooks):
+        return {}
+
+
+class _RecordingModel:
+    def __init__(self):
+        self.calls = []
+
+    def apply_model(self, input_x, timestep, **conditioning):
+        transformer_options = conditioning["transformer_options"]
+        crossattn = conditioning["c_crossattn"]
+        self.calls.append((crossattn.shape[1], timestep.clone(), transformer_options))
+        markers = crossattn[:, 0, 0].reshape(-1, 1, 1, 1)
+        return torch.ones_like(input_x) * markers
+
+
+def test_euler_and_cfg_match_reference_formulas():
+    x = torch.tensor([[[4.0, 2.0]]])
+    denoised = torch.tensor([[[1.0, 0.5]]])
+    sigma = torch.tensor(2.0)
+    sigma_next = torch.tensor(0.75)
+    assert torch.equal(euler_step(x, denoised, sigma, sigma_next), x + (x - denoised) / sigma * (sigma_next - sigma))
+
+    cond = torch.tensor([3.0])
+    uncond = torch.tensor([1.0])
+    assert torch.equal(cfg_combine(cond, uncond, 5.0), torch.tensor([11.0]))
+    assert cfg_combine(cond, uncond, 1.0) is cond
+    assert _cfg_branches(1.0, {}) == (("positive", 0),)
+    assert _cfg_branches(1.0, {"disable_cfg1_optimization": True}) == (("negative", 1), ("positive", 0))
+    assert _cfg_branches(5.0, {}) == (("negative", 1), ("positive", 0))
+
+
+def test_single_request_prediction_uses_standard_sampling_function(monkeypatch):
+    expected = torch.ones(1, 2)
+    seen = []
+
+    def sampling_function(model, x, sigma, negative, positive, cfg, model_options, seed):
+        seen.append((model, x, sigma, negative, positive, cfg, model_options, seed))
+        return expected
+
+    monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.sampling_function", sampling_function)
+    session = ContinuousBatchSession(object())
+    session.inner_model = "inner-model"
+    session.model_options = {"transformer_options": {"sample_sigmas": torch.tensor([1.0, 0.0])}}
+    state = SimpleNamespace(
+        x=torch.zeros(1, 2),
+        sigmas=torch.tensor([2.0, 0.0]),
+        index=0,
+        conds={"negative": ["negative"], "positive": ["positive"]},
+        cfg=5.0,
+        seed=42,
+    )
+
+    assert session.predict([state]) == [expected]
+    assert seen[0][0] == "inner-model"
+    assert torch.equal(seen[0][2], torch.tensor([2.0]))
+    assert seen[0][3:6] == (["negative"], ["positive"], 5.0)
+    assert torch.equal(seen[0][6]["transformer_options"]["sample_sigmas"], state.sigmas)
+    assert torch.equal(session.model_options["transformer_options"]["sample_sigmas"], torch.tensor([1.0, 0.0]))
+    assert seen[0][7] == 42
+
+
+def test_multi_prediction_buckets_positive_154_and_negative_77_for_two_requests(monkeypatch):
+    monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.get_area_and_mult", lambda cond, *args: cond)
+    patcher = _RecordingPatcher()
+    model = _RecordingModel()
+    session = ContinuousBatchSession(patcher)
+    session.inner_model = model
+    session.model_options = {}
+    states = [
+        _batch_state(2.0, 2.0, 1.0, 3.0),
+        _batch_state(1.0, 3.0, 10.0, 20.0),
+    ]
+
+    session.predict(states)
+
+    assert [call[0] for call in model.calls] == [77, 154]
+    assert torch.equal(patcher.prepared_sigmas[0], torch.tensor([2.0, 1.0]))
+    for _, sigmas, transformer_options in model.calls:
+        assert torch.equal(sigmas, torch.tensor([2.0, 1.0]))
+        assert torch.equal(transformer_options["sigmas"], sigmas)
+    assert model.calls[0][2]["cond_or_uncond"] == [1, 1]
+    assert model.calls[1][2]["cond_or_uncond"] == [0, 0]
+    assert model.calls[0][2]["uuids"] == ["negative-1.0", "negative-10.0"]
+    assert model.calls[1][2]["uuids"] == ["positive-3.0", "positive-20.0"]
+
+
+def test_multi_prediction_remaps_bucket_outputs_before_cfg(monkeypatch):
+    monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.get_area_and_mult", lambda cond, *args: cond)
+    session = ContinuousBatchSession(_RecordingPatcher())
+    session.inner_model = _RecordingModel()
+    session.model_options = {}
+    states = [
+        _batch_state(2.0, 2.0, 1.0, 3.0),
+        _batch_state(1.0, 3.0, 10.0, 20.0),
+    ]
+
+    predictions = session.predict(states)
+
+    assert torch.equal(predictions[0], torch.full_like(states[0].x, 5.0))
+    assert torch.equal(predictions[1], torch.full_like(states[1].x, 40.0))
+
+
+def test_model_family_validation_accepts_only_plain_sd_contracts():
+    sd15 = _bare_model(comfy.model_base.BaseModel, comfy.supported_models.SD15)
+    _validate_model_family(FAMILY_SD15, _model_patcher(sd15))
+
+    sd20 = _bare_model(comfy.model_base.BaseModel, comfy.supported_models.SD20)
+    with pytest.raises(ValueError, match="standard SD1.5"):
+        _validate_model_family(FAMILY_SD15, _model_patcher(sd20))
+
+    inpaint = _bare_model(comfy.model_base.BaseModel, comfy.supported_models.SD15, in_channels=9)
+    with pytest.raises(ValueError, match="inpaint or concatenated"):
+        _validate_model_family(FAMILY_SD15, _model_patcher(inpaint))
+
+    sdxl = _bare_model(comfy.model_base.SDXL, comfy.supported_models.SDXL)
+    refiner = _bare_model(comfy.model_base.SDXLRefiner, comfy.supported_models.SDXLRefiner)
+    _validate_model_family(FAMILY_SDXL, _model_patcher(sdxl))
+    _validate_model_family(FAMILY_SDXL, _model_patcher(refiner))
+
+    ip2p = _bare_model(comfy.model_base.SDXL_instructpix2pix, comfy.supported_models.SDXL_instructpix2pix, in_channels=8)
+    with pytest.raises(ValueError, match="standard SDXL"):
+        _validate_model_family(FAMILY_SDXL, _model_patcher(ip2p))
+
+
+def test_model_extension_validation_is_conservative_for_sd_and_anima():
+    for family in (FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL):
+        _validate_model_extensions(family, _extension_patcher())
+    with pytest.raises(ValueError, match="model callbacks"):
+        _validate_model_extensions(FAMILY_ANIMA, _extension_patcher(callbacks={"event": {None: [object()]}}))
+    with pytest.raises(ValueError, match="additional models"):
+        _validate_model_extensions(FAMILY_SD15, _extension_patcher(additional_models={"control": [object()]}))
+
+
+@pytest.mark.parametrize("family", [FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL])
+@pytest.mark.parametrize("patch_key", ["patches", "patches_replace"])
+def test_model_extension_validation_rejects_transformer_patches_for_every_family(family, patch_key):
+    with pytest.raises(ValueError, match="transformer patches"):
+        _validate_model_extensions(family, _extension_patcher(model_options={"transformer_options": {patch_key: {"attn": [object()]}}}))
+
+
+def test_model_extension_validation_allows_anima_attention_backend_override():
+    _validate_model_extensions(FAMILY_ANIMA, _extension_patcher(model_options={"transformer_options": {"optimized_attention_override": object()}}))
+
+
+@pytest.mark.parametrize("family", [FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL])
+@pytest.mark.parametrize("wrapper_type", [comfy.patcher_extension.WrappersMP.APPLY_MODEL, comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL])
+def test_model_extension_validation_rejects_wrappers_for_every_family(family, wrapper_type):
+    wrappers = {wrapper_type: {None: [object()]}}
+    with pytest.raises(ValueError, match="wrapper"):
+        _validate_model_extensions(family, _extension_patcher(wrappers=wrappers))
+
+
+def test_conditioning_contracts_reject_unsupported_features_and_wrappers():
+    raw = [[torch.zeros(1, 77, 768), {}]]
+    _validate_conditioning("positive", raw)
+    with pytest.raises(ValueError, match="one positive"):
+        _validate_conditioning("positive", raw + raw)
+    with pytest.raises(ValueError, match="control"):
+        _validate_conditioning("positive", [[raw[0][0], {"control": object()}]])
+
+    sd15 = _processed_cond({"c_crossattn": comfy.conds.CONDCrossAttn(torch.zeros(1, 77, 768))})
+    _processed_conditioning_signature(FAMILY_SD15, sd15)
+    with pytest.raises(ValueError, match="unsupported c_crossattn"):
+        _processed_conditioning_signature(FAMILY_SD15, _processed_cond({"c_crossattn": comfy.conds.CONDRegular(torch.zeros(1, 77, 768))}))
+
+    sdxl = _processed_cond({
+        "c_crossattn": comfy.conds.CONDCrossAttn(torch.zeros(1, 77, 2048)),
+        "y": comfy.conds.CONDRegular(torch.zeros(1, 2816)),
+    })
+    _processed_conditioning_signature(FAMILY_SDXL, sdxl)
+    with pytest.raises(ValueError, match="conditioning keys"):
+        _processed_conditioning_signature(FAMILY_SDXL, sd15)
+
+
+def test_single_request_rejects_unsupported_processed_conditioning_during_prepare(monkeypatch):
+    bad_condition = _processed_cond({"c_crossattn": comfy.conds.CONDRegular(torch.zeros(1, 77, 768))})
+    monkeypatch.setattr("comfy.continuous_batching.comfy.sampler_helpers.convert_cond", lambda cond: cond)
+    monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.process_conds", lambda *args, **kwargs: {"positive": [bad_condition], "negative": [bad_condition]})
+    monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.get_area_and_mult", lambda cond, *args: cond)
+
+    model_sampling = SimpleNamespace(
+        sigma_max=torch.tensor(1.0),
+        noise_scaling=lambda sigma, noise, latent, max_denoise: noise,
+    )
+    session = ContinuousBatchSession(SimpleNamespace(load_device=torch.device("cpu")))
+    session.inner_model = SimpleNamespace(model_sampling=model_sampling)
+    state = SimpleNamespace(
+        family=FAMILY_SD15,
+        noise=torch.zeros(1, 4, 2, 2),
+        latent_image=torch.zeros(1, 4, 2, 2),
+        sigmas=torch.tensor([1.0, 0.0]),
+        positive=[object()],
+        negative=[object()],
+        seed=1,
+        prepared=False,
+    )
+
+    with pytest.raises(ValueError, match="unsupported c_crossattn conditioning wrapper"):
+        session.prepare_request(state)
+    assert not state.prepared
+
+
+def test_anima_conditioning_structure_groups_padding_compatible_prompts():
+    short = [[torch.zeros(1, 12, 1024), {"t5xxl_ids": torch.zeros(1, 77)}]]
+    longer = [[torch.zeros(1, 40, 1024), {"t5xxl_ids": torch.zeros(1, 200)}]]
+    long = [[torch.zeros(1, 12, 1024), {"t5xxl_ids": torch.zeros(1, 513)}]]
+    assert _conditioning_structure(FAMILY_ANIMA, short) == _conditioning_structure(FAMILY_ANIMA, longer)
+    assert _conditioning_structure(FAMILY_ANIMA, short) != _conditioning_structure(FAMILY_ANIMA, long)
+
+
+def test_callback_uses_request_progress_and_routing_context():
+    request_registry = ProgressRegistry("request", DynamicPrompt({}))
+    reset_progress_state("outer", DynamicPrompt({}))
+    client_token = set_current_client_id("outer-client")
+    seen = []
+    state = SimpleNamespace(
+        callback=lambda *args: seen.append((get_executing_context(), get_current_client_id(), get_progress_state())),
+        client_id="request-client",
+        progress_registry=request_registry,
+        prompt_id="request",
+        node_id="sampler",
+        index=0,
+        x=torch.zeros(1),
+        sigmas=torch.tensor([1.0, 0.0]),
+    )
+    try:
+        ContinuousBatchSession.run_callback(state, torch.zeros(1))
+        assert seen[0][0].prompt_id == "request"
+        assert seen[0][0].node_id == "sampler"
+        assert seen[0][1] == "request-client"
+        assert seen[0][2] is request_registry
+        assert get_current_client_id() == "outer-client"
+        assert get_progress_state().prompt_id == "outer"
+    finally:
+        reset_current_client_id(client_token)
+
+
+def test_admits_at_step_boundaries_and_retires_finished_requests():
+    async def run():
+        first = FakeState("first", 3)
+        coordinator = ContinuousBatchCoordinator("key", first)
+        second = FakeState("second", 1)
+        second_tasks = []
+        session = FakeSession(on_first_step=lambda: second_tasks.append(asyncio.create_task(coordinator.submit(second))))
+        coordinator.session = session
+        first_task = asyncio.create_task(coordinator.submit(first))
+        for _ in range(100):
+            if second_tasks:
+                break
+            if first_task.done():
+                await first_task
+            await asyncio.sleep(0)
+        assert second_tasks
+        assert await asyncio.gather(first_task, second_tasks[0]) == ["first output", "second output"]
+        assert session.batches[0] == ["first"]
+        assert ["first", "second"] in session.batches
+        assert session.batches[-1] == ["first"]
+        assert first.cleared and second.cleared
+        assert session.closed == 1
+
+    asyncio.run(run())
+
+
+def test_failure_clears_requests_and_finishing_request_reloads_survivor():
+    async def failure():
+        first = FakeState("first", 1)
+        second = FakeState("second", 1)
+        coordinator = ContinuousBatchCoordinator("key", first)
+        session = FakeSession(fail=True)
+        coordinator.session = session
+        results = await asyncio.gather(coordinator.submit(first), coordinator.submit(second), return_exceptions=True)
+        assert all(isinstance(result, RuntimeError) for result in results)
+        assert first.cleared and second.cleared
+        assert session.closed == 1
+
+    async def residency():
+        fast = FakeState("fast", 1)
+        slow = FakeState("slow", 2)
+        coordinator = ContinuousBatchCoordinator("key", fast)
+        session = FakeSession()
+        coordinator.session = session
+        assert await asyncio.gather(coordinator.submit(fast), coordinator.submit(slow)) == ["fast output", "slow output"]
+        assert session.batches == [["fast", "slow"], ["slow"]]
+        assert session.reload_requests == 1
+
+    asyncio.run(failure())
+    asyncio.run(residency())

tests-unit/comfy_test/test_continuous_batching_cli.py

@@ -0,0 +1,60 @@
+import subprocess
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[2]
+IMPORT_CLI = "import comfy.options; comfy.options.enable_args_parsing(); import comfy.cli_args"
+PRINT_CONTINUOUS_BATCHING = IMPORT_CLI + "; print(comfy.cli_args.args.continuous_batching)"
+PRINT_DYNAMIC_VRAM = IMPORT_CLI + "; print(comfy.cli_args.enables_dynamic_vram())"
+
+
+def run_cli(*arguments, script=IMPORT_CLI):
+    return subprocess.run(
+        [sys.executable, "-c", script, *arguments],
+        cwd=ROOT,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+
+def test_continuous_batching_defaults_to_disabled():
+    result = run_cli(script=PRINT_CONTINUOUS_BATCHING)
+    assert result.returncode == 0
+    assert result.stdout.strip() == "0"
+
+
+def test_continuous_batching_stores_the_prompt_limit():
+    result = run_cli("--continuous-batching", "4", script=PRINT_CONTINUOUS_BATCHING)
+    assert result.returncode == 0
+    assert result.stdout.strip() == "4"
+
+
+def test_continuous_batching_rejects_negative_limit():
+    result = run_cli("--continuous-batching", "-1")
+    assert result.returncode == 2
+    assert "must be zero or greater" in result.stderr
+
+
+def test_multi_prompt_continuous_batching_rejects_cache_none():
+    result = run_cli("--continuous-batching", "2", "--cache-none")
+    assert result.returncode == 2
+    assert "incompatible with --cache-none" in result.stderr
+
+
+def test_single_prompt_continuous_batching_allows_cache_none():
+    result = run_cli("--continuous-batching", "1", "--cache-none")
+    assert result.returncode == 0
+
+
+def test_continuous_batching_disables_automatic_dynamic_vram():
+    result = run_cli("--continuous-batching", "1", script=PRINT_DYNAMIC_VRAM)
+    assert result.returncode == 0
+    assert result.stdout.strip() == "False"
+
+
+def test_continuous_batching_rejects_explicit_dynamic_vram():
+    result = run_cli("--continuous-batching", "2", "--enable-dynamic-vram")
+    assert result.returncode == 2
+    assert "incompatible with --enable-dynamic-vram" in result.stderr

tests-unit/execution_test/test_continuous_worker.py

@@ -0,0 +1,152 @@
+import asyncio
+
+import torch
+
+import execution
+import main
+
+
+class StopWorker(Exception):
+    pass
+
+
+def test_cooperative_worker_enters_owned_inference_scope(monkeypatch):
+    observed = []
+
+    class Queue:
+        def set_cooperative(self, enabled):
+            pass
+
+        def is_cancelled(self, prompt_id):
+            return False
+
+        def get_if(self, predicate):
+            observed.append(torch.is_inference_mode_enabled())
+            raise StopWorker()
+
+    monkeypatch.setattr(main, "prompt_executor_config", lambda: (execution.CacheType.NONE, {"lru": 0, "ram": 0, "ram_inactive": 0}))
+    monkeypatch.setattr(main.comfy.memory_management, "set_ram_cache_release_state", lambda *args: None)
+    monkeypatch.setattr(main.comfy.continuous_batching, "set_cancel_checker", lambda checker: None)
+
+    async def run():
+        try:
+            await main.cooperative_prompt_worker(Queue(), object(), 2)
+        except StopWorker:
+            pass
+
+    asyncio.run(run())
+    assert observed == [True]
+    assert not torch.is_inference_mode_enabled()
+
+
+def test_cooperative_group_owner_clears_interrupt_once_before_start(monkeypatch):
+    interrupt_calls = []
+
+    class Queue:
+        def __init__(self):
+            self.get_calls = 0
+            self.finished_drains = 0
+
+        def set_cooperative(self, enabled):
+            pass
+
+        def is_cancelled(self, prompt_id):
+            return False
+
+        def get_if(self, predicate):
+            self.get_calls += 1
+            if self.get_calls == 1:
+                return ((0, "prompt", {}, {}, [], {}), 0)
+            raise StopWorker()
+
+        def finish_cooperative_drain(self):
+            self.finished_drains += 1
+
+        def get_flags(self):
+            return {}
+
+    async def execute_prompt(*args, **kwargs):
+        return None
+
+    queue = Queue()
+    monkeypatch.setattr(main, "prompt_executor_config", lambda: (execution.CacheType.NONE, {"lru": 0, "ram": 0, "ram_inactive": 0}))
+    monkeypatch.setattr(main, "execute_prompt_async", execute_prompt)
+    monkeypatch.setattr(main.nodes, "interrupt_processing", lambda value=True: interrupt_calls.append(value))
+    monkeypatch.setattr(main.comfy.memory_management, "set_ram_cache_release_state", lambda *args: None)
+    monkeypatch.setattr(main.comfy.continuous_batching, "set_cancel_checker", lambda checker: None)
+    monkeypatch.setattr(main.comfy.model_management, "soft_empty_cache", lambda: None)
+    monkeypatch.setattr(main.hook_breaker_ac10a0, "restore_functions", lambda: None)
+    monkeypatch.setattr(main.gc, "collect", lambda: None)
+    monkeypatch.setattr(main.asset_seeder, "is_disabled", lambda: True)
+    monkeypatch.setattr(main.asset_seeder, "resume", lambda: None)
+
+    async def run():
+        try:
+            await main.cooperative_prompt_worker(queue, object(), 1)
+        except StopWorker:
+            pass
+
+    asyncio.run(run())
+    assert interrupt_calls == [False]
+    assert queue.finished_drains == 1
+
+
+def test_prompt_worker_preserves_requested_max_prompts(monkeypatch):
+    observed = []
+
+    async def cooperative_worker(q, server_instance, max_prompts):
+        observed.append(max_prompts)
+
+    monkeypatch.setattr(main, "cooperative_prompt_worker", cooperative_worker)
+    for max_prompts in (1, 3):
+        monkeypatch.setattr(main.args, "continuous_batching", max_prompts)
+        main.prompt_worker(object(), object())
+
+    assert observed == [1, 3]
+
+
+def _queue_item(prompt, outputs, preview_method=None):
+    return (0, "prompt-id", prompt, {"preview_method": preview_method}, outputs)
+
+
+def _continuous_prompt(sampler_ids, sampler_type="AnimaContinuousKSampler"):
+    prompt = {
+        "model": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}},
+        "output": {"class_type": "SaveImage", "inputs": {}},
+    }
+    for sampler_id in sampler_ids:
+        prompt[sampler_id] = {"class_type": sampler_type, "inputs": {"model": ["model", 0]}}
+    return prompt
+
+
+def test_continuous_prompt_key_accepts_all_connected_sampler_families():
+    keys = []
+    for sampler_type in main.comfy.continuous_batching.CONTINUOUS_SAMPLER_NODE_FAMILIES:
+        prompt = _continuous_prompt(["sampler"], sampler_type)
+        prompt["output"]["inputs"]["images"] = ["sampler", 0]
+        key = main.continuous_prompt_key(_queue_item(prompt, ["output"]))
+        assert key is not None
+        assert key[0] == sampler_type
+        keys.append(key)
+    assert len(set(keys)) == 3
+
+
+def test_continuous_prompt_key_requires_exactly_one_connected_sampler():
+    disconnected = _continuous_prompt(["sampler"])
+    assert main.continuous_prompt_key(_queue_item(disconnected, ["output"])) is None
+
+    multiple = _continuous_prompt(["sampler-a", "sampler-b"])
+    multiple["output"]["inputs"].update({"first": ["sampler-a", 0], "second": ["sampler-b", 0]})
+    assert main.continuous_prompt_key(_queue_item(multiple, ["output"])) is None
+
+
+def test_continuous_prompt_key_tracks_model_graph_and_preview_method():
+    first = _continuous_prompt(["sampler"])
+    first["output"]["inputs"]["images"] = ["sampler", 0]
+    second = _continuous_prompt(["sampler"])
+    second["output"]["inputs"]["images"] = ["sampler", 0]
+    second["model"]["inputs"]["ckpt_name"] = "other.safetensors"
+
+    first_key = main.continuous_prompt_key(_queue_item(first, ["output"], "auto"))
+    assert first_key != main.continuous_prompt_key(_queue_item(first, ["output"], "none"))
+    assert first_key != main.continuous_prompt_key(_queue_item(second, ["output"], "auto"))

tests-unit/execution_test/test_graph_builder_context.py

@@ -0,0 +1,34 @@
+import asyncio
+
+from comfy_execution.graph_utils import GraphBuilder
+
+
+def test_default_prefix_is_isolated_between_async_tasks():
+    async def run():
+        ready = 0
+        both_ready = asyncio.Event()
+
+        async def build(root, call_index, graph_index):
+            nonlocal ready
+            GraphBuilder.set_default_prefix(root, call_index, graph_index)
+            ready += 1
+            if ready == 2:
+                both_ready.set()
+            await both_ready.wait()
+            return GraphBuilder().prefix, GraphBuilder().prefix
+
+        return await asyncio.gather(build("first", 10, 20), build("second", 30, 40))
+
+    assert asyncio.run(run()) == [
+        ("first.10.20.", "first.10.21."),
+        ("second.30.40.", "second.30.41."),
+    ]
+
+
+def test_explicit_prefix_still_advances_default_graph_index():
+    try:
+        GraphBuilder.set_default_prefix("default", 1, 2)
+        assert GraphBuilder.alloc_prefix("explicit", 3, 4) == "explicit.3.4."
+        assert GraphBuilder.alloc_prefix() == "default.1.3."
+    finally:
+        GraphBuilder.set_default_prefix("", 0, 0)

tests-unit/execution_test/test_prompt_context_isolation.py

@@ -0,0 +1,262 @@
+import asyncio
+from types import SimpleNamespace
+
+import torch
+
+import nodes
+from comfy_execution.caching import BasicCache, CacheKeySetID
+from comfy_execution.graph import DynamicPrompt
+from comfy_execution.progress import WebUIProgressHandler, get_progress_state, reset_progress_state
+from comfy_execution.utils import get_current_client_id, has_current_client_id, reset_current_client_id, set_current_client_id
+from execution import CacheSet, PromptExecutor, _send_cached_ui
+
+
+class Server:
+    def __init__(self):
+        self.client_id = "shared-server-value"
+        self.messages = []
+        self.last_node_id = None
+
+    def send_sync(self, event, data, client_id):
+        self.messages.append((event, data, client_id))
+
+
+def test_prompt_executor_messages_use_executor_client_id():
+    server = Server()
+    first = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+    second = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+    first.client_id = "first-client"
+    second.client_id = "second-client"
+
+    first.add_message("event", {"prompt_id": "first"}, broadcast=False)
+    second.add_message("event", {"prompt_id": "second"}, broadcast=False)
+
+    assert [message[2] for message in server.messages] == ["first-client", "second-client"]
+    assert server.client_id == "shared-server-value"
+
+
+def test_cached_ui_is_recorded_without_a_connected_client():
+    server = Server()
+    ui_outputs = {}
+    cached = SimpleNamespace(ui={"output": {"images": []}, "meta": {"node_id": "node"}})
+
+    _send_cached_ui(server, None, "node", "node", cached, "prompt", ui_outputs)
+
+    assert ui_outputs == {"node": cached.ui}
+    assert server.messages == []
+
+
+def test_progress_registry_is_isolated_between_async_prompt_tasks():
+    async def prompt(prompt_id, ready, release):
+        reset_progress_state(prompt_id, DynamicPrompt({}))
+        ready.set()
+        await release.wait()
+        return get_progress_state().prompt_id
+
+    async def run():
+        first_ready = asyncio.Event()
+        second_ready = asyncio.Event()
+        release = asyncio.Event()
+        first = asyncio.create_task(prompt("first", first_ready, release))
+        second = asyncio.create_task(prompt("second", second_ready, release))
+        await asyncio.gather(first_ready.wait(), second_ready.wait())
+        release.set()
+        return await asyncio.gather(first, second)
+
+    assert asyncio.run(run()) == ["first", "second"]
+
+
+def test_webui_progress_handler_uses_prompt_client_id():
+    server = Server()
+    first = WebUIProgressHandler(server, "first-client")
+    second = WebUIProgressHandler(server, "second-client")
+
+    first._send_progress_state("first", {})
+    second._send_progress_state("second", {})
+
+    assert [message[2] for message in server.messages] == ["first-client", "second-client"]
+
+
+def test_explicit_anonymous_client_does_not_fall_back_to_stale_server_client():
+    server = Server()
+    handler = WebUIProgressHandler(server, None)
+    handler._send_progress_state("anonymous", {})
+    assert server.messages[0][2] is None
+
+    assert not has_current_client_id()
+    token = set_current_client_id(None)
+    try:
+        assert has_current_client_id()
+        assert get_current_client_id() is None
+    finally:
+        reset_current_client_id(token)
+
+
+def test_shared_cache_uses_task_local_prompt_signatures():
+    async def run():
+        cache = BasicCache(CacheKeySetID)
+        both_ready = asyncio.Event()
+        ready_count = 0
+
+        async def prompt(class_type, value):
+            nonlocal ready_count
+            dynprompt = DynamicPrompt({"same-id": {"class_type": class_type, "inputs": {}}})
+            await cache.set_prompt(dynprompt, ["same-id"], None)
+            cache.set_local("same-id", value)
+            ready_count += 1
+            if ready_count == 2:
+                both_ready.set()
+            await both_ready.wait()
+            result = cache.get_local("same-id")
+            cache.release_prompt()
+            return result
+
+        return await asyncio.gather(prompt("First", "first"), prompt("Second", "second"))
+
+    assert asyncio.run(run()) == ["first", "second"]
+
+
+def test_cooperative_executors_do_not_share_stateful_node_instances():
+    instances = []
+    both_ready = None
+    ready = 0
+
+    class StatefulProbe:
+        def __init__(self):
+            instances.append(self)
+
+        @classmethod
+        def INPUT_TYPES(cls):
+            return {"required": {}}
+
+        RETURN_TYPES = ()
+        FUNCTION = "run"
+
+        async def run(self):
+            nonlocal ready
+            ready += 1
+            if ready == 2:
+                both_ready.set()
+            await both_ready.wait()
+            return ()
+
+    async def run():
+        nonlocal both_ready
+        both_ready = asyncio.Event()
+        server = Server()
+        server.prompt_queue = SimpleNamespace(cooperative=True, is_cancelled=lambda prompt_id: False)
+        cache_args = {"lru": 0, "ram": 0, "ram_inactive": 0}
+        shared_outputs = CacheSet(cache_args=cache_args).outputs
+        prompt = {"same-id": {"class_type": "StatefulInterleaveProbe", "inputs": {}}}
+        first = PromptExecutor(server, cache_args=cache_args, shared_outputs=shared_outputs)
+        second = PromptExecutor(server, cache_args=cache_args, shared_outputs=shared_outputs)
+        with torch.inference_mode():
+            await asyncio.gather(
+                first.execute_async(prompt, "first", {}, ["same-id"]),
+                second.execute_async(prompt, "second", {}, ["same-id"]),
+            )
+
+    nodes.NODE_CLASS_MAPPINGS["StatefulInterleaveProbe"] = StatefulProbe
+    try:
+        asyncio.run(run())
+    finally:
+        del nodes.NODE_CLASS_MAPPINGS["StatefulInterleaveProbe"]
+    assert len(instances) == 2
+    assert instances[0] is not instances[1]
+
+
+def test_concurrent_cooperative_executors_do_not_clear_global_interrupt(monkeypatch):
+    interrupt_calls = []
+    monkeypatch.setattr(nodes, "interrupt_processing", lambda value=True: interrupt_calls.append(value))
+
+    async def run():
+        server = Server()
+        server.prompt_queue = SimpleNamespace(cooperative=True, is_cancelled=lambda prompt_id: False)
+        first = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+        second = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+        with torch.inference_mode():
+            await asyncio.gather(
+                first.execute_async({}, "first", {}, []),
+                second.execute_async({}, "second", {}, []),
+            )
+
+    asyncio.run(run())
+    assert interrupt_calls == []
+
+
+def test_inference_mode_stays_enabled_across_interleaved_cooperative_tasks():
+    observations = []
+    both_ready = None
+    ready = 0
+
+    class Probe:
+        @classmethod
+        def INPUT_TYPES(cls):
+            return {"required": {}}
+
+        RETURN_TYPES = ()
+        FUNCTION = "run"
+
+        async def run(self):
+            nonlocal ready
+            observations.append(torch.is_inference_mode_enabled())
+            ready += 1
+            if ready == 2:
+                both_ready.set()
+            await both_ready.wait()
+            observations.append(torch.is_inference_mode_enabled())
+            return ()
+
+    async def run():
+        nonlocal both_ready
+        both_ready = asyncio.Event()
+        server = Server()
+        server.prompt_queue = SimpleNamespace(cooperative=True, is_cancelled=lambda prompt_id: False)
+        prompt = {"probe": {"class_type": "InferenceModeInterleaveProbe", "inputs": {}}}
+        first = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+        second = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+        with torch.inference_mode():
+            await asyncio.gather(
+                first.execute_async(prompt, "first", {}, ["probe"]),
+                second.execute_async(prompt, "second", {}, ["probe"]),
+            )
+            assert torch.is_inference_mode_enabled()
+
+    nodes.NODE_CLASS_MAPPINGS["InferenceModeInterleaveProbe"] = Probe
+    try:
+        asyncio.run(run())
+    finally:
+        del nodes.NODE_CLASS_MAPPINGS["InferenceModeInterleaveProbe"]
+    assert observations == [True, True, True, True]
+    assert not torch.is_inference_mode_enabled()
+
+
+def test_serial_prompt_executor_still_owns_inference_mode():
+    observations = []
+
+    class Probe:
+        @classmethod
+        def INPUT_TYPES(cls):
+            return {"required": {}}
+
+        RETURN_TYPES = ()
+        FUNCTION = "run"
+
+        def run(self):
+            observations.append(torch.is_inference_mode_enabled())
+            return ()
+
+    async def run():
+        server = Server()
+        server.prompt_queue = SimpleNamespace(cooperative=False)
+        prompt = {"probe": {"class_type": "InferenceModeSerialProbe", "inputs": {}}}
+        executor = PromptExecutor(server, cache_args={"lru": 0, "ram": 0, "ram_inactive": 0})
+        await executor.execute_async(prompt, "serial", {}, ["probe"])
+
+    nodes.NODE_CLASS_MAPPINGS["InferenceModeSerialProbe"] = Probe
+    try:
+        asyncio.run(run())
+    finally:
+        del nodes.NODE_CLASS_MAPPINGS["InferenceModeSerialProbe"]
+    assert observations == [True]
+    assert not torch.is_inference_mode_enabled()

tests/execution/test_prompt_queue.py

@@ -0,0 +1,117 @@
+import execution
+
+
+class TestServer:
+    def queue_updated(self):
+        pass
+
+
+def test_cooperative_cancel_is_prompt_scoped_without_node_class_checks(monkeypatch):
+    interrupted = []
+    monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
+    queue = execution.PromptQueue(TestServer())
+    prompt = {"1": {"class_type": "UnrelatedNode", "inputs": {}}}
+    queue.put((0, "first", prompt, {}, [], {}))
+    queue.put((1, "second", prompt, {}, [], {}))
+    _, first_id = queue.get(timeout=0)
+    _, second_id = queue.get(timeout=0)
+
+    queue.set_cooperative(True)
+    assert queue.interrupt_if_running("first")
+    assert queue.is_cancelled("first")
+    assert not queue.is_cancelled("second")
+    assert interrupted == []
+
+    queue.task_done(first_id, {}, None)
+    assert not queue.is_cancelled("first")
+    queue.task_done(second_id, {}, None)
+
+
+def test_single_cooperative_cancel_also_wakes_synchronous_execution(monkeypatch):
+    interrupted = []
+    monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
+    queue = execution.PromptQueue(TestServer())
+    queue.put((0, "only", {}, {}, [], {}))
+    _, item_id = queue.get(timeout=0)
+    queue.set_cooperative(True)
+
+    assert queue.interrupt_if_running("only")
+    assert queue.is_cancelled("only")
+    assert queue.is_cooperative_draining()
+    assert interrupted == [True]
+
+    queue.task_done(item_id, {}, None)
+    queue.finish_cooperative_drain()
+
+
+def test_completed_result_wins_if_cancel_arrives_after_compute(monkeypatch):
+    monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: None)
+    queue = execution.PromptQueue(TestServer())
+    queue.put((0, "completed", {}, {}, [], {}))
+    _, item_id = queue.get(timeout=0)
+    queue.set_cooperative(True)
+
+    assert queue.interrupt_if_running("completed")
+    status = execution.PromptQueue.ExecutionStatus("success", True, [])
+    queue.task_done(item_id, {"outputs": {"node": "result"}}, status)
+
+    assert not queue.is_cancelled("completed")
+    history = queue.get_history("completed")["completed"]
+    assert history["status"]["status_str"] == "success"
+    assert history["outputs"] == {"node": "result"}
+
+
+def test_global_cooperative_interrupt_cancels_all_and_blocks_admission_until_drain(monkeypatch):
+    interrupted = []
+    monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
+    queue = execution.PromptQueue(TestServer())
+    for number, prompt_id in enumerate(("first", "second", "next")):
+        queue.put((number, prompt_id, {}, {}, [], {}))
+    _, first_id = queue.get(timeout=0)
+    _, second_id = queue.get(timeout=0)
+    queue.set_cooperative(True)
+
+    queue.interrupt_all_running()
+
+    assert queue.is_cancelled("first")
+    assert queue.is_cancelled("second")
+    assert interrupted == [True]
+    assert queue.get_if(lambda item: True) is None
+
+    queue.task_done(first_id, {}, None)
+    queue.task_done(second_id, {}, None)
+    queue.finish_cooperative_drain()
+    item, next_id = queue.get_if(lambda item: True)
+    assert item[1] == "next"
+    queue.task_done(next_id, {}, None)
+
+
+def test_legacy_cancel_uses_global_interrupt(monkeypatch):
+    interrupted = []
+    monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
+    queue = execution.PromptQueue(TestServer())
+    queue.put((0, "prompt", {}, {}, [], {}))
+    queue.get(timeout=0)
+
+    assert queue.interrupt_if_running("prompt")
+    assert interrupted == [True]
+    assert not queue.is_cancelled("prompt")
+
+
+def test_targeted_cancel_does_not_interrupt_different_serial_prompt(monkeypatch):
+    interrupted = []
+    monkeypatch.setattr(execution.nodes, "interrupt_processing", lambda: interrupted.append(True))
+    queue = execution.PromptQueue(TestServer())
+    queue.put((0, "running", {}, {}, [], {}))
+    queue.get(timeout=0)
+
+    assert not queue.interrupt_if_running("different")
+    assert interrupted == []
+
+
+def test_get_if_leaves_incompatible_head_queued():
+    queue = execution.PromptQueue(TestServer())
+    queue.put((0, "first", {}, {}, [], {}))
+
+    assert queue.get_if(lambda item: item[1] == "other") is None
+    assert queue.get(timeout=0)[0][1] == "first"