Reuse continuous batch conditioning

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))
+                entries.append((state_index, name, branch, state.x, sigma, state.processed_conds[name]))
-                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]):
+                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([entry[5].conditioning 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(
@@ -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)

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}")],
+            "negative": [negative],
-            "positive": [_batch_cond(x, 154, positive_marker, f"positive-{positive_marker}")],
+            "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):
+def test_multi_prediction_remaps_bucket_outputs_before_cfg():
-    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 = {}
@@ -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))