diff --git a/__init__.py b/__init__.py index c749b76..7d64d96 100644 --- a/__init__.py +++ b/__init__.py @@ -30,8 +30,8 @@ from .simple_syrup.runtime.attention_region_prompt_handler import ( # noqa: E40 from .simple_syrup.runtime.comfy_safetensors_dtypes import ( # noqa: E402 register_comfy_safetensors_dtypes, ) -from .simple_syrup.runtime.noise_reuse_prompt_handler import ( # noqa: E402 - register_noise_reuse_prompt_handler, +from .simple_syrup.runtime.noise_reuse_validation import ( # noqa: E402 + register_noise_reuse_validation, ) WEB_DIRECTORY = "./web/dist" @@ -61,7 +61,7 @@ register_quant_cache_routes() register_external_llm_routes() register_mask_batch_preview_routes() register_attention_region_prompt_handler() -register_noise_reuse_prompt_handler() +register_noise_reuse_validation() __all__ = [ "WEB_DIRECTORY", diff --git a/simple_syrup/runtime/noise_reuse_prompt_handler.py b/simple_syrup/runtime/noise_reuse_prompt_handler.py deleted file mode 100644 index df52a18..0000000 --- a/simple_syrup/runtime/noise_reuse_prompt_handler.py +++ /dev/null @@ -1,40 +0,0 @@ -# SimpleSyrup - workflow-focused ComfyUI extensions for image generation -# Copyright (C) 2026 Artificial Sweetener and contributors -# SPDX-License-Identifier: AGPL-3.0-or-later -"""Register capture coordination before Comfy validates or executes source sampling.""" - -from __future__ import annotations - -from importlib import import_module -from typing import Any - -from ..shared.logging import get_logger -from .noise_reuse_graph import rewrite_noise_reuse_graph - -LOGGER = get_logger(__name__) -_MARKER = "_simple_syrup_noise_reuse_prompt_handler" - - -def register_noise_reuse_prompt_handler() -> None: - """Register once when the native PromptServer becomes available.""" - try: - instance = getattr(import_module("server").PromptServer, "instance", None) - except (ImportError, AttributeError): - return - if instance is not None and not getattr(instance, _MARKER, False): - instance.add_on_prompt_handler(_rewrite_prompt) - 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): - rewrite_noise_reuse_graph( - data["prompt"], import_module("nodes").NODE_CLASS_MAPPINGS - ) - return data diff --git a/simple_syrup/runtime/noise_reuse_validation.py b/simple_syrup/runtime/noise_reuse_validation.py new file mode 100644 index 0000000..c23db34 --- /dev/null +++ b/simple_syrup/runtime/noise_reuse_validation.py @@ -0,0 +1,87 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Prepare capture at Comfy's native validation boundary for every queue caller.""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from functools import wraps +from importlib import import_module +from typing import Any, Protocol, cast + +from ..shared.logging import get_logger +from .noise_reuse_graph import rewrite_noise_reuse_graph + +LOGGER = get_logger(__name__) +_MARKER = "_simple_syrup_noise_reuse_validation" +NativePromptValidator = Callable[ + [str, dict[str, Any], list[str] | None], + Awaitable[tuple[bool, object, object, object]], +] + + +class NativeValidationModule(Protocol): + """Describe the host's replaceable async prompt validation boundary.""" + + validate_prompt: NativePromptValidator + + +def prepare_noise_reuse_prompt(prompt: dict[str, Any]) -> None: + """Install cache-keyed capture dependencies before source sampling executes.""" + rewrite_noise_reuse_graph(prompt, import_module("nodes").NODE_CLASS_MAPPINGS) + + +class NoiseReuseValidationAdapter: + """Own capture preparation while preserving the host validator's contract.""" + + def __init__(self, validator: NativePromptValidator) -> None: + """Retain the current validator, including other installed extensions.""" + self._validator = validator + + def wrap(self) -> NativePromptValidator: + """Preserve the host's async function identity and introspection metadata.""" + + @wraps(self._validator) + async def validate_prompt( + prompt_id: str, + prompt: dict[str, Any], + partial_execution_list: list[str] | None, + ) -> tuple[bool, object, object, object]: + """Delegate through this adapter without changing the native signature.""" + return await self(prompt_id, prompt, partial_execution_list) + + setattr(validate_prompt, _MARKER, True) + return validate_prompt + + async def __call__( + self, + prompt_id: str, + prompt: dict[str, Any], + partial_execution_list: list[str] | None, + ) -> tuple[bool, object, object, object]: + """Prepare once before native validation and propagate its exact result.""" + try: + prepare_noise_reuse_prompt(prompt) + except Exception: + LOGGER.exception( + "Noise reuse capture preparation failed before Comfy validation", + extra={"prompt_id": prompt_id, "operation": "prepare_capture"}, + ) + raise + return await self._validator(prompt_id, prompt, partial_execution_list) + + +def register_noise_reuse_validation() -> None: + """Install once on the shared boundary used by HTTP and direct graph queues.""" + try: + execution = cast(NativeValidationModule, import_module("execution")) + except ModuleNotFoundError as error: + if error.name != "execution": + raise + return + validator = execution.validate_prompt + if not getattr(validator, _MARKER, False): + execution.validate_prompt = NoiseReuseValidationAdapter(validator).wrap() + LOGGER.info("Installed noise reuse capture preparation before Comfy validation") diff --git a/simple_syrup/services/noise_reuse_graph_preparation_service.py b/simple_syrup/services/noise_reuse_graph_preparation_service.py index 9f7b337..14ebae1 100644 --- a/simple_syrup/services/noise_reuse_graph_preparation_service.py +++ b/simple_syrup/services/noise_reuse_graph_preparation_service.py @@ -1,16 +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.""" +"""Expose capture preparation through the existing public Cube capability.""" from __future__ import annotations -from ..runtime.noise_reuse_prompt_handler import prepare_cube_prompt +from ..runtime.noise_reuse_validation import prepare_noise_reuse_prompt class NoiseReuseGraphPreparationService: - """Share the capture graph authority with queues outside Comfy's prompt hook.""" + """Delegate public host requests to the authoritative capture planner.""" def prepare(self, prompt: dict[str, dict[str, object]]) -> None: """Install cache-keyed source capture before any sampler executes.""" - prepare_cube_prompt(prompt) + prepare_noise_reuse_prompt(prompt) diff --git a/tests/comfy_integration/test_noise_reuse_validation.py b/tests/comfy_integration/test_noise_reuse_validation.py new file mode 100644 index 0000000..0889fe8 --- /dev/null +++ b/tests/comfy_integration/test_noise_reuse_validation.py @@ -0,0 +1,259 @@ +# SimpleSyrup - workflow-focused ComfyUI extensions for image generation +# Copyright (C) 2026 Artificial Sweetener and contributors +# SPDX-License-Identifier: AGPL-3.0-or-later + +"""Prove capture preparation precedes native validation for every queue caller.""" + +from __future__ import annotations + +import asyncio +import inspect +import sys +from copy import deepcopy +from types import ModuleType +from typing import Any + +import pytest + +from simple_syrup.runtime.noise_reuse_validation import ( + NativePromptValidator, + prepare_noise_reuse_prompt, + register_noise_reuse_validation, +) + + +class ValidationModule(ModuleType): + """Describe the isolated host module's replaceable validation function.""" + + validate_prompt: NativePromptValidator + + +async def validate( + execution: ValidationModule, + prompt_id: str, + prompt: dict[str, Any], + partial_targets: list[str] | None, +) -> tuple[bool, object, object, object]: + """Await the external validator whether it returns a coroutine or awaitable.""" + return await execution.validate_prompt(prompt_id, prompt, partial_targets) + + +def capture_graph() -> dict[str, Any]: + """Describe direct source lineage and a live receiving capture dependency.""" + return { + "source": { + "class_type": "KSampler", + "inputs": { + "model": ["loader", 0], + "latent_image": ["empty", 0], + "sampler_name": "euler_ancestral", + }, + }, + "reuse": { + "class_type": "SimpleSyrup.NoiseReuseOptions", + "inputs": {"enabled": True}, + }, + "receive": { + "class_type": "SimpleSyrup.KSampler", + "inputs": { + "model": ["loader", 0], + "latent_image": ["source", 0], + "options": ["reuse", 0], + "steps": 12, + "sampler_name": "euler_ancestral", + "scheduler": "normal", + "denoise": 0.15, + }, + }, + } + + +@pytest.fixture +def validation_host( + monkeypatch: pytest.MonkeyPatch, +) -> tuple[ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]]]: + """Isolate only Comfy's external validation and registration boundaries.""" + calls: list[tuple[str, dict[str, Any], list[str] | None]] = [] + + async def validate_prompt( + prompt_id: str, + prompt: dict[str, Any], + partial_execution_list: list[str] | None, + ) -> tuple[bool, object, object, object]: + """Record the graph exactly as the native validator receives it.""" + calls.append((prompt_id, deepcopy(prompt), partial_execution_list)) + return True, None, ["receive"], {} + + execution = ValidationModule("execution") + execution.validate_prompt = validate_prompt + nodes = ModuleType("nodes") + nodes.NODE_CLASS_MAPPINGS = {} # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "execution", execution) + monkeypatch.setitem(sys.modules, "nodes", nodes) + return execution, calls + + +@pytest.mark.parametrize("partial_targets", [None, ["receive"]]) +def test_direct_validation_prepares_capture_without_http_handler( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], + partial_targets: list[str] | None, +) -> None: + """Cover both full Cube queueing and Comfy partial-output validation.""" + execution, calls = validation_host + register_noise_reuse_validation() + prompt = capture_graph() + result = asyncio.run(validate(execution, "run", prompt, partial_targets)) + assert result == (True, None, ["receive"], {}) + assert len(calls) == 1 + prompt_id, validated, targets = calls[0] + assert prompt_id == "run" and targets is partial_targets + assert validated == prompt + assert validated["source"]["inputs"]["model"] == ["source.trajectory_capture", 0] + assert validated["receive"]["inputs"]["latent_image"] == [ + "source.sampling_provenance", + 0, + ] + assert validated["receive.noise_capture_request"]["inputs"]["denoise"] == 0.15 + + +def test_registration_and_history_requeue_do_not_duplicate_capture( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], +) -> None: + """Preserve graph identity across registration, pre-preparation and history.""" + execution, calls = validation_host + register_noise_reuse_validation() + validator = execution.validate_prompt + register_noise_reuse_validation() + assert execution.validate_prompt is validator + prompt = capture_graph() + prepare_noise_reuse_prompt(prompt) + prepared = deepcopy(prompt) + for prompt_id in ("first", "requeue"): + asyncio.run(validate(execution, prompt_id, prompt, None)) + assert prompt == prepared and len(calls) == 2 + assert all(validated == prepared for _, validated, _ in calls) + + +def test_registration_preserves_async_signature_and_wrapped_validator( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], +) -> None: + """Keep native function inspection and extension chaining compatible.""" + execution, _ = validation_host + native = execution.validate_prompt + register_noise_reuse_validation() + wrapped = execution.validate_prompt + assert inspect.iscoroutinefunction(wrapped) + assert inspect.signature(wrapped) == inspect.signature(native) + assert inspect.unwrap(wrapped) is native + + +def test_changed_receiving_controls_remain_source_cache_dependencies( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], +) -> None: + """Keep fresh submissions sensitive to downstream denoise and live steps.""" + execution, calls = validation_host + register_noise_reuse_validation() + for denoise in (0.15, 0.2): + prompt = capture_graph() + prompt["receive"]["inputs"].update(denoise=denoise, steps=["live_steps", 0]) + asyncio.run(validate(execution, "run", prompt, None)) + requests = [ + graph["receive.noise_capture_request"]["inputs"] for _, graph, _ in calls + ] + assert [request["denoise"] for request in requests] == [0.15, 0.2] + assert all(request["steps"] == ["live_steps", 0] for request in requests) + + +def test_unrelated_graph_and_validation_rejection_are_preserved( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], +) -> None: + """Leave ordinary workflows and native validation error payloads intact.""" + execution, _ = validation_host + prompt = {"output": {"class_type": "SaveImage", "inputs": {}}} + before = deepcopy(prompt) + rejection: tuple[bool, object, object, object] = ( + False, + {"type": "native_error"}, + [], + {"output": {"errors": []}}, + ) + + async def reject( + prompt_id: str, prompt: dict[str, Any], partial_execution_list: list[str] | None + ) -> tuple[bool, object, object, object]: + """Return the host's validation rejection without extension translation.""" + return rejection + + execution.validate_prompt = reject + register_noise_reuse_validation() + assert asyncio.run(validate(execution, "run", prompt, None)) is rejection + assert prompt == before + + +def test_native_validation_exception_propagates( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], +) -> None: + """Keep unexpected host failures visible to the queue that called it.""" + execution, _ = validation_host + + async def fail( + prompt_id: str, prompt: dict[str, Any], partial_execution_list: list[str] | None + ) -> tuple[bool, object, object, object]: + """Represent a failed external native validator.""" + raise RuntimeError("native validation failed") + + execution.validate_prompt = fail + register_noise_reuse_validation() + with pytest.raises(RuntimeError, match="native validation failed"): + asyncio.run(validate(execution, "run", capture_graph(), None)) + + +def test_preparation_failure_stops_before_validation( + validation_host: tuple[ + ValidationModule, list[tuple[str, dict[str, Any], list[str] | None]] + ], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Reject an unexpected planner failure rather than silently skipping capture.""" + execution, calls = validation_host + + def fail(prompt: dict[str, Any]) -> None: + """Represent failure at the graph preparation boundary.""" + raise RuntimeError("capture planning failed") + + monkeypatch.setattr( + "simple_syrup.runtime.noise_reuse_validation.prepare_noise_reuse_prompt", fail + ) + register_noise_reuse_validation() + with pytest.raises(RuntimeError, match="capture planning failed"): + asyncio.run(validate(execution, "run", capture_graph(), None)) + assert calls == [] + assert "capture planning failed" in caplog.text + + +def test_registration_without_comfy_runtime_is_safe( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Allow tooling imports when Comfy's execution module is unavailable.""" + + def unavailable(name: str) -> ModuleType: + """Represent an absent external host installation.""" + raise ModuleNotFoundError("No module named execution", name="execution") + + monkeypatch.setattr( + "simple_syrup.runtime.noise_reuse_validation.import_module", unavailable + ) + register_noise_reuse_validation()