From d0ba49e2c921b41e625dddf32c362f5479c786fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=87=8E=E7=94=9F=E3=81=AE=E7=94=B7?= Date: Fri, 17 Jul 2026 02:39:55 +0900 Subject: [PATCH] Reuse continuous batch conditioning --- comfy/continuous_batching.py | 33 +++++++-- .../comfy_test/test_continuous_batching.py | 73 +++++++++++++++++-- 2 files changed, 93 insertions(+), 13 deletions(-) diff --git a/comfy/continuous_batching.py b/comfy/continuous_batching.py index 07c4d3286..bde1041dd 100644 --- a/comfy/continuous_batching.py +++ b/comfy/continuous_batching.py @@ -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) diff --git a/tests-unit/comfy_test/test_continuous_batching.py b/tests-unit/comfy_test/test_continuous_batching.py index 5140fc397..5ca15ef0f 100644 --- a/tests-unit/comfy_test/test_continuous_batching.py +++ b/tests-unit/comfy_test/test_continuous_batching.py @@ -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))