import asyncio from types import SimpleNamespace import pytest import torch import comfy.conds import comfy.model_base import comfy.patcher_extension import comfy.supported_models from comfy.continuous_batching import ( FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL, ContinuousBatchCoordinator, ContinuousBatchSession, _cfg_branches, _conditioning_structure, _PreparedConditioning, _processed_conditioning_signature, _validate_conditioning, _validate_model_extensions, _validate_model_family, cfg_combine, euler_step, ) from comfy_execution.graph import DynamicPrompt from comfy_execution.progress import ProgressRegistry, get_progress_state, reset_progress_state from comfy_execution.utils import get_current_client_id, get_executing_context, reset_current_client_id, set_current_client_id class FakeState: def __init__(self, name, steps, max_batch_size=2, admission_delay=0.0): self.name = name self.steps = steps self.max_batch_size = max_batch_size self.admission_delay = admission_delay self.model_patcher = object() self.family = "test" self.output = None self.cleared = False def clear(self): self.cleared = True def key(self): return "group" class FakeSession: def __init__(self, fail=False, on_first_step=None): self.batches = [] self.closed = 0 self.open = False self.fail = fail self.on_first_step = on_first_step self.reload_requests = 0 def step(self, states): self.open = True self.batches.append([state.name for state in states]) if len(self.batches) == 1 and self.on_first_step is not None: self.on_first_step() if self.fail: raise RuntimeError("denoiser failed") updates = [] for state in states: state.steps -= 1 finished = state.steps == 0 if finished: state.output = state.name + " output" updates.append((state, finished)) return updates def close(self): self.closed += 1 self.open = False def request_model_reload(self): self.reload_requests += 1 def _model_patcher(model): return SimpleNamespace(model=model) def _bare_model(model_type, config_type, in_channels=4, latent_channels=4, concat_keys=()): model = object.__new__(model_type) torch.nn.Module.__init__(model) model.model_config = object.__new__(config_type) model.diffusion_model = SimpleNamespace(in_channels=in_channels) model.latent_format = SimpleNamespace(latent_channels=latent_channels) model.concat_keys = concat_keys return model def _extension_patcher(model_options=None, callbacks=None, additional_models=None, wrappers=None): return SimpleNamespace( model_options=model_options or {}, callbacks=callbacks or {}, additional_models=additional_models or {}, wrappers=wrappers or {}, ) def _processed_cond(conditioning): return SimpleNamespace(area=None, control=None, patches=None, hooks=None, conditioning=conditioning) def _batch_cond(x, length, marker, uuid): cond = _processed_cond({"c_crossattn": comfy.conds.CONDCrossAttn(torch.full((1, length, 4), marker))}) cond.input_x = x cond.uuid = uuid return cond 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, sigmas=torch.tensor([sigma, 0.0]), index=0, cfg=cfg, conds={ "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)), }, ) class _RecordingPatcher: def __init__(self): self.prepared_sigmas = [] def prepare_state(self, sigmas, model_options): self.prepared_sigmas.append(sigmas.clone()) def apply_hooks(self, hooks): return {} class _RecordingModel: def __init__(self): self.calls = [] def apply_model(self, input_x, timestep, **conditioning): transformer_options = conditioning["transformer_options"] crossattn = conditioning["c_crossattn"] self.calls.append((crossattn.shape[1], timestep.clone(), transformer_options)) markers = crossattn[:, 0, 0].reshape(-1, 1, 1, 1) return torch.ones_like(input_x) * markers def test_euler_and_cfg_match_reference_formulas(): x = torch.tensor([[[4.0, 2.0]]]) denoised = torch.tensor([[[1.0, 0.5]]]) sigma = torch.tensor(2.0) sigma_next = torch.tensor(0.75) assert torch.equal(euler_step(x, denoised, sigma, sigma_next), x + (x - denoised) / sigma * (sigma_next - sigma)) cond = torch.tensor([3.0]) uncond = torch.tensor([1.0]) assert torch.equal(cfg_combine(cond, uncond, 5.0), torch.tensor([11.0])) assert cfg_combine(cond, uncond, 1.0) is cond assert _cfg_branches(1.0, {}) == (("positive", 0),) assert _cfg_branches(1.0, {"disable_cfg1_optimization": True}) == (("negative", 1), ("positive", 0)) assert _cfg_branches(5.0, {}) == (("negative", 1), ("positive", 0)) def test_single_request_prediction_uses_standard_sampling_function(monkeypatch): expected = torch.ones(1, 2) seen = [] def sampling_function(model, x, sigma, negative, positive, cfg, model_options, seed): seen.append((model, x, sigma, negative, positive, cfg, model_options, seed)) return expected monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.sampling_function", sampling_function) session = ContinuousBatchSession(object()) session.inner_model = "inner-model" session.model_options = {"transformer_options": {"sample_sigmas": torch.tensor([1.0, 0.0])}} state = SimpleNamespace( x=torch.zeros(1, 2), sigmas=torch.tensor([2.0, 0.0]), index=0, conds={"negative": ["negative"], "positive": ["positive"]}, cfg=5.0, seed=42, ) assert session.predict([state]) == [expected] assert seen[0][0] == "inner-model" assert torch.equal(seen[0][2], torch.tensor([2.0])) assert seen[0][3:6] == (["negative"], ["positive"], 5.0) assert torch.equal(seen[0][6]["transformer_options"]["sample_sigmas"], state.sigmas) assert torch.equal(session.model_options["transformer_options"]["sample_sigmas"], torch.tensor([1.0, 0.0])) assert seen[0][7] == 42 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 *args: pytest.fail("predict reprocessed conditioning")) patcher = _RecordingPatcher() model = _RecordingModel() session = ContinuousBatchSession(patcher) session.inner_model = model session.model_options = {} states = [ _batch_state(2.0, 2.0, 1.0, 3.0), _batch_state(1.0, 3.0, 10.0, 20.0), ] session.predict(states) assert [call[0] for call in model.calls] == [77, 154] assert torch.equal(patcher.prepared_sigmas[0], torch.tensor([2.0, 1.0])) for _, sigmas, transformer_options in model.calls: assert torch.equal(sigmas, torch.tensor([2.0, 1.0])) assert torch.equal(transformer_options["sigmas"], sigmas) assert model.calls[0][2]["cond_or_uncond"] == [1, 1] assert model.calls[1][2]["cond_or_uncond"] == [0, 0] assert model.calls[0][2]["uuids"] == ["negative-1.0", "negative-10.0"] assert model.calls[1][2]["uuids"] == ["positive-3.0", "positive-20.0"] def test_multi_prediction_remaps_bucket_outputs_before_cfg(): session = ContinuousBatchSession(_RecordingPatcher()) session.inner_model = _RecordingModel() session.model_options = {} states = [ _batch_state(2.0, 2.0, 1.0, 3.0), _batch_state(1.0, 3.0, 10.0, 20.0), ] predictions = session.predict(states) assert torch.equal(predictions[0], torch.full_like(states[0].x, 5.0)) 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)) sd20 = _bare_model(comfy.model_base.BaseModel, comfy.supported_models.SD20) with pytest.raises(ValueError, match="standard SD1.5"): _validate_model_family(FAMILY_SD15, _model_patcher(sd20)) inpaint = _bare_model(comfy.model_base.BaseModel, comfy.supported_models.SD15, in_channels=9) with pytest.raises(ValueError, match="inpaint or concatenated"): _validate_model_family(FAMILY_SD15, _model_patcher(inpaint)) sdxl = _bare_model(comfy.model_base.SDXL, comfy.supported_models.SDXL) refiner = _bare_model(comfy.model_base.SDXLRefiner, comfy.supported_models.SDXLRefiner) _validate_model_family(FAMILY_SDXL, _model_patcher(sdxl)) _validate_model_family(FAMILY_SDXL, _model_patcher(refiner)) ip2p = _bare_model(comfy.model_base.SDXL_instructpix2pix, comfy.supported_models.SDXL_instructpix2pix, in_channels=8) with pytest.raises(ValueError, match="standard SDXL"): _validate_model_family(FAMILY_SDXL, _model_patcher(ip2p)) def test_model_extension_validation_is_conservative_for_sd_and_anima(): for family in (FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL): _validate_model_extensions(family, _extension_patcher()) with pytest.raises(ValueError, match="model callbacks"): _validate_model_extensions(FAMILY_ANIMA, _extension_patcher(callbacks={"event": {None: [object()]}})) with pytest.raises(ValueError, match="additional models"): _validate_model_extensions(FAMILY_SD15, _extension_patcher(additional_models={"control": [object()]})) @pytest.mark.parametrize("family", [FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL]) @pytest.mark.parametrize("patch_key", ["patches", "patches_replace"]) def test_model_extension_validation_rejects_transformer_patches_for_every_family(family, patch_key): with pytest.raises(ValueError, match="transformer patches"): _validate_model_extensions(family, _extension_patcher(model_options={"transformer_options": {patch_key: {"attn": [object()]}}})) def test_model_extension_validation_allows_anima_attention_backend_override(): _validate_model_extensions(FAMILY_ANIMA, _extension_patcher(model_options={"transformer_options": {"optimized_attention_override": object()}})) @pytest.mark.parametrize("family", [FAMILY_ANIMA, FAMILY_SD15, FAMILY_SDXL]) @pytest.mark.parametrize("wrapper_type", [comfy.patcher_extension.WrappersMP.APPLY_MODEL, comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL]) def test_model_extension_validation_rejects_wrappers_for_every_family(family, wrapper_type): wrappers = {wrapper_type: {None: [object()]}} with pytest.raises(ValueError, match="wrapper"): _validate_model_extensions(family, _extension_patcher(wrappers=wrappers)) def test_conditioning_contracts_reject_unsupported_features_and_wrappers(): raw = [[torch.zeros(1, 77, 768), {}]] _validate_conditioning("positive", raw) with pytest.raises(ValueError, match="one positive"): _validate_conditioning("positive", raw + raw) with pytest.raises(ValueError, match="control"): _validate_conditioning("positive", [[raw[0][0], {"control": object()}]]) sd15 = _processed_cond({"c_crossattn": comfy.conds.CONDCrossAttn(torch.zeros(1, 77, 768))}) _processed_conditioning_signature(FAMILY_SD15, sd15) with pytest.raises(ValueError, match="unsupported c_crossattn"): _processed_conditioning_signature(FAMILY_SD15, _processed_cond({"c_crossattn": comfy.conds.CONDRegular(torch.zeros(1, 77, 768))})) sdxl = _processed_cond({ "c_crossattn": comfy.conds.CONDCrossAttn(torch.zeros(1, 77, 2048)), "y": comfy.conds.CONDRegular(torch.zeros(1, 2816)), }) _processed_conditioning_signature(FAMILY_SDXL, sdxl) with pytest.raises(ValueError, match="conditioning keys"): _processed_conditioning_signature(FAMILY_SDXL, sd15) def test_single_request_rejects_unsupported_processed_conditioning_during_prepare(monkeypatch): bad_condition = _processed_cond({"c_crossattn": comfy.conds.CONDRegular(torch.zeros(1, 77, 768))}) monkeypatch.setattr("comfy.continuous_batching.comfy.sampler_helpers.convert_cond", lambda cond: cond) monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.process_conds", lambda *args, **kwargs: {"positive": [bad_condition], "negative": [bad_condition]}) monkeypatch.setattr("comfy.continuous_batching.comfy.samplers.get_area_and_mult", lambda cond, *args: cond) model_sampling = SimpleNamespace( sigma_max=torch.tensor(1.0), noise_scaling=lambda sigma, noise, latent, max_denoise: noise, ) session = ContinuousBatchSession(SimpleNamespace(load_device=torch.device("cpu"))) session.inner_model = SimpleNamespace(model_sampling=model_sampling) state = SimpleNamespace( family=FAMILY_SD15, noise=torch.zeros(1, 4, 2, 2), latent_image=torch.zeros(1, 4, 2, 2), sigmas=torch.tensor([1.0, 0.0]), positive=[object()], negative=[object()], seed=1, prepared=False, ) with pytest.raises(ValueError, match="unsupported c_crossattn conditioning wrapper"): session.prepare_request(state) assert not state.prepared def test_anima_conditioning_structure_groups_padding_compatible_prompts(): short = [[torch.zeros(1, 12, 1024), {"t5xxl_ids": torch.zeros(1, 77)}]] longer = [[torch.zeros(1, 40, 1024), {"t5xxl_ids": torch.zeros(1, 200)}]] long = [[torch.zeros(1, 12, 1024), {"t5xxl_ids": torch.zeros(1, 513)}]] assert _conditioning_structure(FAMILY_ANIMA, short) == _conditioning_structure(FAMILY_ANIMA, longer) assert _conditioning_structure(FAMILY_ANIMA, short) != _conditioning_structure(FAMILY_ANIMA, long) def test_callback_uses_request_progress_and_routing_context(): request_registry = ProgressRegistry("request", DynamicPrompt({})) reset_progress_state("outer", DynamicPrompt({})) client_token = set_current_client_id("outer-client") seen = [] state = SimpleNamespace( callback=lambda *args: seen.append((get_executing_context(), get_current_client_id(), get_progress_state())), client_id="request-client", progress_registry=request_registry, prompt_id="request", node_id="sampler", index=0, x=torch.zeros(1), sigmas=torch.tensor([1.0, 0.0]), ) try: ContinuousBatchSession.run_callback(state, torch.zeros(1)) assert seen[0][0].prompt_id == "request" assert seen[0][0].node_id == "sampler" assert seen[0][1] == "request-client" assert seen[0][2] is request_registry assert get_current_client_id() == "outer-client" assert get_progress_state().prompt_id == "outer" finally: reset_current_client_id(client_token) def test_admits_at_step_boundaries_and_retires_finished_requests(): async def run(): first = FakeState("first", 3) coordinator = ContinuousBatchCoordinator("key", first) second = FakeState("second", 1) second_tasks = [] session = FakeSession(on_first_step=lambda: second_tasks.append(asyncio.create_task(coordinator.submit(second)))) coordinator.session = session first_task = asyncio.create_task(coordinator.submit(first)) for _ in range(100): if second_tasks: break if first_task.done(): await first_task await asyncio.sleep(0) assert second_tasks assert await asyncio.gather(first_task, second_tasks[0]) == ["first output", "second output"] assert session.batches[0] == ["first"] assert ["first", "second"] in session.batches assert session.batches[-1] == ["first"] assert first.cleared and second.cleared assert session.closed == 1 asyncio.run(run()) def test_failure_clears_requests_and_finishing_request_reloads_survivor(): async def failure(): first = FakeState("first", 1) second = FakeState("second", 1) coordinator = ContinuousBatchCoordinator("key", first) session = FakeSession(fail=True) coordinator.session = session results = await asyncio.gather(coordinator.submit(first), coordinator.submit(second), return_exceptions=True) assert all(isinstance(result, RuntimeError) for result in results) assert first.cleared and second.cleared assert session.closed == 1 async def residency(): fast = FakeState("fast", 1) slow = FakeState("slow", 2) coordinator = ContinuousBatchCoordinator("key", fast) session = FakeSession() coordinator.session = session assert await asyncio.gather(coordinator.submit(fast), coordinator.submit(slow)) == ["fast output", "slow output"] assert session.batches == [["fast", "slow"], ["slow"]] assert session.reload_requests == 1 asyncio.run(failure()) asyncio.run(residency())