diff --git a/simple_syrup/nodes_v3/ksampler.py b/simple_syrup/nodes_v3/ksampler.py index 3890d53..d043a86 100644 --- a/simple_syrup/nodes_v3/ksampler.py +++ b/simple_syrup/nodes_v3/ksampler.py @@ -10,6 +10,9 @@ from typing import Any, ClassVar from ..domain.sampler_options import SamplerOptions from ..nodes import tooltips +from ..services.noise_reuse_graph_preparation_service import ( + NoiseReuseGraphPreparationService, +) from ..services.sampler_options_sampling_service import SamplerOptionsSamplingService from .ksampler_schema import ksampler_inputs from .sampler_options_schema import COMFY_IO, OptionsNodeBase, options_input @@ -21,6 +24,12 @@ class KSamplerV3(OptionsNodeBase): service_class: ClassVar[type[SamplerOptionsSamplingService]] = ( SamplerOptionsSamplingService ) + SUGARCUBES_PROMPT_PREPARATION_API_VERSION: ClassVar[int] = 1 + + @classmethod + def prepare_sugarcubes_prompt(cls, prompt: dict[str, dict[str, object]]) -> None: + """Install source capture dependencies before Cube queue validation.""" + NoiseReuseGraphPreparationService().prepare(prompt) @classmethod def define_schema(cls) -> Any: diff --git a/simple_syrup/runtime/noise_reuse_prompt_handler.py b/simple_syrup/runtime/noise_reuse_prompt_handler.py index 0deaf83..df52a18 100644 --- a/simple_syrup/runtime/noise_reuse_prompt_handler.py +++ b/simple_syrup/runtime/noise_reuse_prompt_handler.py @@ -26,6 +26,11 @@ def register_noise_reuse_prompt_handler() -> None: setattr(instance, _MARKER, True) +def prepare_cube_prompt(prompt: dict[str, dict[str, object]]) -> None: + """Reuse the authoritative capture planner on the final lowered Cube graph.""" + rewrite_noise_reuse_graph(prompt, import_module("nodes").NODE_CLASS_MAPPINGS) + + def _rewrite_prompt(data: Any) -> Any: """Use normal error-producing nodes for invalid source provenance.""" if isinstance(data, dict) and isinstance(data.get("prompt"), dict): diff --git a/simple_syrup/services/noise_reuse_graph_preparation_service.py b/simple_syrup/services/noise_reuse_graph_preparation_service.py new file mode 100644 index 0000000..9f7b337 --- /dev/null +++ b/simple_syrup/services/noise_reuse_graph_preparation_service.py @@ -0,0 +1,16 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later +"""Prepare captured-noise dependencies for direct Cube queue execution.""" + +from __future__ import annotations + +from ..runtime.noise_reuse_prompt_handler import prepare_cube_prompt + + +class NoiseReuseGraphPreparationService: + """Share the capture graph authority with queues outside Comfy's prompt hook.""" + + def prepare(self, prompt: dict[str, dict[str, object]]) -> None: + """Install cache-keyed source capture before any sampler executes.""" + prepare_cube_prompt(prompt) diff --git a/tests/comfy_integration/test_noise_reuse_graph.py b/tests/comfy_integration/test_noise_reuse_graph.py index 56f0d5b..bb45ff9 100644 --- a/tests/comfy_integration/test_noise_reuse_graph.py +++ b/tests/comfy_integration/test_noise_reuse_graph.py @@ -4,11 +4,14 @@ """Prove receiver planning, strict image provenance and source cache dependencies.""" +import sys from copy import deepcopy +from types import ModuleType from typing import Any import pytest +from simple_syrup.nodes_v3.ksampler import KSamplerV3 from simple_syrup.runtime.noise_reuse_graph import ( rewrite_noise_reuse_graph, source_lineage, @@ -57,6 +60,27 @@ def graph() -> dict[str, Any]: } +def test_cube_queue_capability_installs_capture_before_source_sampling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Share exact capture planning with direct Cube queues that omit prompt hooks.""" + registry_module = ModuleType("nodes") + monkeypatch.setattr(registry_module, "NODE_CLASS_MAPPINGS", {}, raising=False) + monkeypatch.setitem(sys.modules, "nodes", registry_module) + prompt = graph() + assert KSamplerV3.SUGARCUBES_PROMPT_PREPARATION_API_VERSION == 1 + KSamplerV3.prepare_sugarcubes_prompt(prompt) + assert prompt["source"]["inputs"]["model"] == ["source.trajectory_capture", 0] + assert prompt["decode"]["inputs"]["samples"] == ["source.sampling_provenance", 0] + assert ( + prompt["source.trajectory_capture"]["class_type"] + == "SimpleSyrup.TrajectoryCaptureModel" + ) + first = deepcopy(prompt) + KSamplerV3.prepare_sugarcubes_prompt(prompt) + assert prompt == first + + def test_image_bridge_resolves_original_latent_and_both_consumers_share_capture() -> ( None ):