mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
663 lines
26 KiB
Python
663 lines
26 KiB
Python
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)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _PreparedConditioning:
|
|
conditioning: dict
|
|
uuid: Any
|
|
signature: tuple
|
|
|
|
|
|
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
|
|
processed_conds: dict | None = field(default=None, init=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.processed_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)
|
|
processed_conds = {}
|
|
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)
|
|
signature = _processed_conditioning_signature(state.family, processed)
|
|
processed_conds[name] = _PreparedConditioning(
|
|
conditioning=processed.conditioning,
|
|
uuid=processed.uuid,
|
|
signature=signature,
|
|
)
|
|
state.latent_image = latent_image
|
|
state.conds = conds
|
|
state.processed_conds = processed_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:
|
|
entries.append((state_index, name, branch, state.x, sigma, state.processed_conds[name]))
|
|
|
|
buckets = []
|
|
for entry in entries:
|
|
for bucket in buckets:
|
|
if (
|
|
entry[5].signature == bucket[0][5].signature
|
|
and entry[3].shape == bucket[0][3].shape
|
|
and comfy.samplers.cond_equal_size(entry[5].conditioning, bucket[0][5].conditioning)
|
|
):
|
|
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])
|
|
conditioning = comfy.samplers.cond_cat([entry[5].conditioning for entry in bucket])
|
|
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"] = [entry[5].uuid for entry in bucket]
|
|
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)
|