diff --git a/comfy/cli_args.py b/comfy/cli_args.py index e2e0d97ec..d0962184b 100644 --- a/comfy/cli_args.py +++ b/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 diff --git a/comfy/continuous_batching.py b/comfy/continuous_batching.py new file mode 100644 index 000000000..07c4d3286 --- /dev/null +++ b/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) diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index f0fdf1aa5..65822cf54 100644 --- a/comfy/sd1_clip.py +++ b/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) diff --git a/comfy/text_encoders/anima.py b/comfy/text_encoders/anima.py index 2e31b2b04..4754b95b5 100644 --- a/comfy/text_encoders/anima.py +++ b/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) diff --git a/comfy/text_encoders/anima_cache.py b/comfy/text_encoders/anima_cache.py new file mode 100644 index 000000000..8392acd17 --- /dev/null +++ b/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 diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index 6bd99b68f..be0b6f8e6 100644 --- a/comfy_execution/caching.py +++ b/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: diff --git a/comfy_execution/graph_utils.py b/comfy_execution/graph_utils.py index 496d2c634..9c69084de 100644 --- a/comfy_execution/graph_utils.py +++ b/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): diff --git a/comfy_execution/progress.py b/comfy_execution/progress.py index 731b8dc66..c0a998476 100644 --- a/comfy_execution/progress.py +++ b/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 diff --git a/comfy_execution/utils.py b/comfy_execution/utils.py index 62d32f101..7824ab1bf 100644 --- a/comfy_execution/utils.py +++ b/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. diff --git a/execution.py b/execution.py index 387772629..5ced90944 100644 --- a/execution.py +++ b/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) diff --git a/main.py b/main.py index 580074b19..cf1bc2a0d 100644 --- a/main.py +++ b/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) diff --git a/nodes.py b/nodes.py index b03d6c603..f5c45db20 100644 --- a/nodes.py +++ b/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)", diff --git a/server.py b/server.py index e28fe2d22..e6156e263 100644 --- a/server.py +++ b/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) diff --git a/tests-unit/comfy_test/test_anima_cache.py b/tests-unit/comfy_test/test_anima_cache.py new file mode 100644 index 000000000..7a3a036b4 --- /dev/null +++ b/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()) == [] diff --git a/tests-unit/comfy_test/test_continuous_batching.py b/tests-unit/comfy_test/test_continuous_batching.py new file mode 100644 index 000000000..5140fc397 --- /dev/null +++ b/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()) diff --git a/tests-unit/comfy_test/test_continuous_batching_cli.py b/tests-unit/comfy_test/test_continuous_batching_cli.py new file mode 100644 index 000000000..287ed27e8 --- /dev/null +++ b/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 diff --git a/tests-unit/execution_test/test_continuous_worker.py b/tests-unit/execution_test/test_continuous_worker.py new file mode 100644 index 000000000..6e1b7be71 --- /dev/null +++ b/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")) diff --git a/tests-unit/execution_test/test_graph_builder_context.py b/tests-unit/execution_test/test_graph_builder_context.py new file mode 100644 index 000000000..04bfc4691 --- /dev/null +++ b/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) diff --git a/tests-unit/execution_test/test_prompt_context_isolation.py b/tests-unit/execution_test/test_prompt_context_isolation.py new file mode 100644 index 000000000..396f1a6ae --- /dev/null +++ b/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() diff --git a/tests/execution/test_prompt_queue.py b/tests/execution/test_prompt_queue.py new file mode 100644 index 000000000..981d5ef73 --- /dev/null +++ b/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"