From 06c5f7d584ba8021ea2594ab25b8225ae2a97666 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 08:43:58 +0900 Subject: [PATCH] Revert "Batch continuous sampler updates" This reverts commit 70465618f2982b9fd8ee2d166aee09799ac9c387. --- comfy/continuous_batching.py | 50 ++------ .../comfy_test/test_continuous_batching.py | 111 ------------------ 2 files changed, 10 insertions(+), 151 deletions(-) diff --git a/comfy/continuous_batching.py b/comfy/continuous_batching.py index ad5092468..bde1041dd 100644 --- a/comfy/continuous_batching.py +++ b/comfy/continuous_batching.py @@ -208,37 +208,6 @@ def cfg_combine(cond, uncond, cfg): return uncond + (cond - uncond) * cfg -def _batch_cfg_predictions(states, state_branches, branch_outputs): - predictions = [None] * len(states) - guided_indices = [] - for index, (state, branches, outputs) in enumerate(zip(states, state_branches, branch_outputs)): - if len(branches) == 1 or math.isclose(state.cfg, 1.0): - predictions[index] = outputs["positive"] - else: - guided_indices.append(index) - - if guided_indices: - cond = torch.cat([branch_outputs[index]["positive"] for index in guided_indices]) - uncond = torch.cat([branch_outputs[index]["negative"] for index in guided_indices]) - cfg = cond.new_tensor([states[index].cfg for index in guided_indices]).reshape( - len(guided_indices), *((1,) * (cond.ndim - 1)) - ) - guided = torch.addcmul(uncond, cond - uncond, cfg).split(1) - for index, prediction in zip(guided_indices, guided): - predictions[index] = prediction - return predictions - - -def _batch_euler_updates(states, denoised): - x = torch.cat([state.x for state in states]) - predictions = torch.cat(denoised) - sigma = torch.stack([state.sigmas[state.index] for state in states]).to(x).reshape( - len(states), *((1,) * (x.ndim - 1)) - ) - sigma_next = torch.stack([state.sigmas[state.index + 1] for state in states]).to(x).reshape_as(sigma) - return torch.addcmul(x, x - predictions, (sigma_next - sigma) / sigma).split(1) - - def _cfg_branches(cfg, model_options): if math.isclose(cfg, 1.0) and not model_options.get("disable_cfg1_optimization", False): return (("positive", 0),) @@ -507,7 +476,13 @@ class ContinuousBatchSession: for entry, output in zip(bucket, outputs): branch_outputs[entry[0]][entry[1]] = output - return _batch_cfg_predictions(states, state_branches, branch_outputs) + 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): @@ -533,18 +508,13 @@ class ContinuousBatchSession: 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) - batched_updates = _batch_euler_updates(states, denoised) if len(states) > 1 else None - updates = [] - for state, prediction, batched_update in zip(states, denoised, batched_updates or [None] * len(states)): - if batched_update is None: - sigma = state.sigmas[state.index].to(state.x) - state.x = euler_step(state.x, prediction, sigma, state.sigmas[state.index + 1].to(state.x)) - else: - state.x = batched_update + 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: diff --git a/tests-unit/comfy_test/test_continuous_batching.py b/tests-unit/comfy_test/test_continuous_batching.py index 19e5ae619..5ca15ef0f 100644 --- a/tests-unit/comfy_test/test_continuous_batching.py +++ b/tests-unit/comfy_test/test_continuous_batching.py @@ -1,5 +1,4 @@ import asyncio -from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -15,8 +14,6 @@ from comfy.continuous_batching import ( FAMILY_SDXL, ContinuousBatchCoordinator, ContinuousBatchSession, - _batch_cfg_predictions, - _batch_euler_updates, _cfg_branches, _conditioning_structure, _PreparedConditioning, @@ -177,114 +174,6 @@ def test_euler_and_cfg_match_reference_formulas(): assert _cfg_branches(5.0, {}) == (("negative", 1), ("positive", 0)) -def test_batched_cfg_and_euler_match_per_request_math_for_different_schedules(monkeypatch): - calls = {"cat": 0, "stack": 0, "addcmul": []} - original_cat = torch.cat - original_stack = torch.stack - original_addcmul = torch.addcmul - - def cat(*args, **kwargs): - calls["cat"] += 1 - return original_cat(*args, **kwargs) - - def stack(*args, **kwargs): - calls["stack"] += 1 - return original_stack(*args, **kwargs) - - def addcmul(input, tensor1, tensor2, **kwargs): - calls["addcmul"].append((tensor2.shape, tensor2.dtype, tensor2.device)) - return original_addcmul(input, tensor1, tensor2, **kwargs) - - monkeypatch.setattr(torch, "cat", cat) - monkeypatch.setattr(torch, "stack", stack) - monkeypatch.setattr(torch, "addcmul", addcmul) - - states = [ - SimpleNamespace(x=torch.tensor([[[[4.0, 2.0]]]]), sigmas=torch.tensor([2.0, 0.5, 0.0]), index=0, cfg=2.0), - SimpleNamespace(x=torch.tensor([[[[3.0, 6.0]]]]), sigmas=torch.tensor([3.0, 1.5, 0.25, 0.0]), index=1, cfg=3.5), - ] - branches = [(("negative", 1), ("positive", 0))] * 2 - outputs = [ - {"negative": torch.tensor([[[[1.0, 0.5]]]]), "positive": torch.tensor([[[[2.0, 1.5]]]])}, - {"negative": torch.tensor([[[[0.5, 1.0]]]]), "positive": torch.tensor([[[[1.5, 2.0]]]])}, - ] - - predictions = _batch_cfg_predictions(states, branches, outputs) - expected_predictions = [cfg_combine(output["positive"], output["negative"], state.cfg) for state, output in zip(states, outputs)] - updates = _batch_euler_updates(states, predictions) - expected_updates = [ - euler_step(state.x, prediction, state.sigmas[state.index], state.sigmas[state.index + 1]) - for state, prediction in zip(states, expected_predictions) - ] - - for actual, expected in zip(predictions, expected_predictions): - torch.testing.assert_close(actual, expected) - for actual, expected in zip(updates, expected_updates): - torch.testing.assert_close(actual, expected) - assert calls == { - "cat": 4, - "stack": 2, - "addcmul": [ - (torch.Size([2, 1, 1, 1]), torch.float32, torch.device("cpu")), - (torch.Size([2, 1, 1, 1]), torch.float32, torch.device("cpu")), - ], - } - - -def test_batched_cfg_preserves_cfg1_positive_output_identity(): - positive = torch.ones(1, 1, 1, 1) - states = [SimpleNamespace(cfg=1.0), SimpleNamespace(cfg=2.0)] - branches = [(("negative", 1), ("positive", 0))] * 2 - outputs = [ - {"negative": torch.zeros_like(positive), "positive": positive}, - {"negative": torch.zeros_like(positive), "positive": torch.full_like(positive, 2.0)}, - ] - - predictions = _batch_cfg_predictions(states, branches, outputs) - - assert predictions[0] is positive - assert torch.equal(predictions[1], torch.full_like(positive, 4.0)) - - -def test_multi_step_runs_callbacks_before_one_batched_euler_update(monkeypatch): - events = [] - session = ContinuousBatchSession(SimpleNamespace(load_device=torch.device("cpu"))) - states = [] - predictions = [] - for index, sigma_next in enumerate((0.5, 0.25)): - x = torch.full((1, 1, 1, 1), 2.0 + index) - states.append(SimpleNamespace( - x=x, - sigmas=torch.tensor([1.0 + index, sigma_next, 0.0]), - index=0, - callback=lambda *args, index=index: events.append(f"callback-{index}"), - client_id=None, - progress_registry=None, - prompt_id=None, - node_id=None, - )) - predictions.append(torch.ones_like(x)) - - monkeypatch.setattr("comfy.continuous_batching.comfy.model_management.cuda_device_context", lambda device: nullcontext()) - monkeypatch.setattr(session, "open_session", lambda requests: None) - monkeypatch.setattr(session, "ensure_model_loaded", lambda requests: None) - monkeypatch.setattr(session, "prepare_request", lambda request: None) - monkeypatch.setattr(session, "predict", lambda requests: predictions) - original_batched_euler = _batch_euler_updates - - def batched_euler(requests, denoised): - events.append("batched-euler") - return original_batched_euler(requests, denoised) - - monkeypatch.setattr("comfy.continuous_batching._batch_euler_updates", batched_euler) - - updates = session.step(states) - - assert events == ["callback-0", "callback-1", "batched-euler"] - assert updates == [(states[0], False), (states[1], False)] - assert [state.index for state in states] == [1, 1] - - def test_single_request_prediction_uses_standard_sampling_function(monkeypatch): expected = torch.ones(1, 2) seen = []