mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
Reuse continuous batch conditioning
This commit is contained in:
野生の男
parent
beef915d9e
commit
d0ba49e2c9
@ -195,6 +195,13 @@ def _processed_conditioning_signature(family, cond):
|
||||
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
|
||||
@ -230,6 +237,7 @@ class ContinuousBatchRequest:
|
||||
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)
|
||||
@ -281,6 +289,7 @@ class ContinuousBatchRequest:
|
||||
self.progress_registry = None
|
||||
self.x = None
|
||||
self.conds = None
|
||||
self.processed_conds = None
|
||||
self.output = None
|
||||
|
||||
|
||||
@ -388,13 +397,20 @@ class ContinuousBatchSession:
|
||||
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)
|
||||
_processed_conditioning_signature(state.family, processed)
|
||||
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):
|
||||
@ -423,14 +439,16 @@ class ContinuousBatchSession:
|
||||
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))
|
||||
entries.append((state_index, name, branch, state.x, sigma, state.processed_conds[name]))
|
||||
|
||||
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]):
|
||||
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:
|
||||
@ -442,8 +460,7 @@ class ContinuousBatchSession:
|
||||
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])
|
||||
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(
|
||||
@ -452,7 +469,7 @@ class ContinuousBatchSession:
|
||||
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["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)
|
||||
|
||||
@ -16,6 +16,7 @@ from comfy.continuous_batching import (
|
||||
ContinuousBatchSession,
|
||||
_cfg_branches,
|
||||
_conditioning_structure,
|
||||
_PreparedConditioning,
|
||||
_processed_conditioning_signature,
|
||||
_validate_conditioning,
|
||||
_validate_model_extensions,
|
||||
@ -115,6 +116,8 @@ def _batch_cond(x, length, marker, uuid):
|
||||
|
||||
def _batch_state(sigma, cfg, negative_marker, positive_marker):
|
||||
x = torch.zeros(1, 4, 2, 2)
|
||||
negative = _batch_cond(x, 77, negative_marker, f"negative-{negative_marker}")
|
||||
positive = _batch_cond(x, 154, positive_marker, f"positive-{positive_marker}")
|
||||
return SimpleNamespace(
|
||||
family=FAMILY_SD15,
|
||||
x=x,
|
||||
@ -122,8 +125,12 @@ def _batch_state(sigma, cfg, negative_marker, positive_marker):
|
||||
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}")],
|
||||
"negative": [negative],
|
||||
"positive": [positive],
|
||||
},
|
||||
processed_conds={
|
||||
"negative": _PreparedConditioning(negative.conditioning, negative.uuid, _processed_conditioning_signature(FAMILY_SD15, negative)),
|
||||
"positive": _PreparedConditioning(positive.conditioning, positive.uuid, _processed_conditioning_signature(FAMILY_SD15, positive)),
|
||||
},
|
||||
)
|
||||
|
||||
@ -198,7 +205,7 @@ def test_single_request_prediction_uses_standard_sampling_function(monkeypatch):
|
||||
|
||||
|
||||
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)
|
||||
monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.get_area_and_mult", lambda *args: pytest.fail("predict reprocessed conditioning"))
|
||||
patcher = _RecordingPatcher()
|
||||
model = _RecordingModel()
|
||||
session = ContinuousBatchSession(patcher)
|
||||
@ -222,8 +229,7 @@ def test_multi_prediction_buckets_positive_154_and_negative_77_for_two_requests(
|
||||
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)
|
||||
def test_multi_prediction_remaps_bucket_outputs_before_cfg():
|
||||
session = ContinuousBatchSession(_RecordingPatcher())
|
||||
session.inner_model = _RecordingModel()
|
||||
session.model_options = {}
|
||||
@ -238,6 +244,63 @@ def test_multi_prediction_remaps_bucket_outputs_before_cfg(monkeypatch):
|
||||
assert torch.equal(predictions[1], torch.full_like(states[1].x, 40.0))
|
||||
|
||||
|
||||
def test_prepare_request_processes_conditioning_once_across_predict_steps(monkeypatch):
|
||||
get_area_calls = []
|
||||
|
||||
def get_area_and_mult(cond, *args):
|
||||
get_area_calls.append(cond.uuid)
|
||||
return cond
|
||||
|
||||
monkeypatch.setattr("comfy.continuous_batching.comfy.sampler_helpers.convert_cond", lambda cond: cond)
|
||||
monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.process_conds", lambda model, noise, conds, *args, **kwargs: conds)
|
||||
monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.get_area_and_mult", get_area_and_mult)
|
||||
|
||||
patcher = _RecordingPatcher()
|
||||
patcher.load_device = torch.device("cpu")
|
||||
model = _RecordingModel()
|
||||
model.model_sampling = SimpleNamespace(
|
||||
sigma_max=torch.tensor(2.0),
|
||||
noise_scaling=lambda sigma, noise, latent, max_denoise: noise,
|
||||
)
|
||||
session = ContinuousBatchSession(patcher)
|
||||
session.inner_model = model
|
||||
session.model_options = {}
|
||||
|
||||
def state(negative_marker, positive_marker):
|
||||
x = torch.zeros(1, 4, 2, 2)
|
||||
negative = _batch_cond(x, 77, negative_marker, f"negative-{negative_marker}")
|
||||
positive = _batch_cond(x, 154, positive_marker, f"positive-{positive_marker}")
|
||||
return SimpleNamespace(
|
||||
family=FAMILY_SD15,
|
||||
noise=x.clone(),
|
||||
latent_image=x.clone(),
|
||||
sigmas=torch.tensor([2.0, 1.0, 0.0]),
|
||||
positive=[positive],
|
||||
negative=[negative],
|
||||
seed=1,
|
||||
cfg=2.0,
|
||||
index=0,
|
||||
prepared=False,
|
||||
processed_conds=None,
|
||||
)
|
||||
|
||||
states = [state(1.0, 3.0), state(10.0, 20.0)]
|
||||
for request in states:
|
||||
session.prepare_request(request)
|
||||
|
||||
assert get_area_calls == ["positive-3.0", "negative-1.0", "positive-20.0", "negative-10.0"]
|
||||
first = session.predict(states)
|
||||
for request in states:
|
||||
request.index = 1
|
||||
second = session.predict(states)
|
||||
|
||||
assert get_area_calls == ["positive-3.0", "negative-1.0", "positive-20.0", "negative-10.0"]
|
||||
assert torch.equal(first[0], torch.full_like(states[0].x, 5.0))
|
||||
assert torch.equal(first[1], torch.full_like(states[1].x, 30.0))
|
||||
assert torch.equal(second[0], first[0])
|
||||
assert torch.equal(second[1], first[1])
|
||||
|
||||
|
||||
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))
|
||||
|
||||
Loading…
Reference in New Issue
Block a user