Artificial Sweetener
f863fcf645
Add VAE encode/decode option nodes across legacy and Comfy v3 exports, with tooltip and registration coverage. Preserve cloned model graph composition for differential diffusion in tiled and regional sampling paths.
215 lines
6.3 KiB
Python
215 lines
6.3 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for shared tiled diffusion sampling mode dispatch."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from simple_syrup.services.tiled_diffusion_sampling_service import (
|
|
TiledDiffusionSamplingService,
|
|
)
|
|
|
|
|
|
def test_multidiffusion_mode_routes_to_multidiffusion_runtime(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""MultiDiffusion mode calls only the MultiDiffusion runtime."""
|
|
|
|
calls: dict[str, dict[str, Any]] = {}
|
|
output: dict[str, Any] = {"samples": torch.ones((1, 4, 4, 4))}
|
|
|
|
def fake_multidiffusion(**kwargs: Any) -> dict[str, Any]:
|
|
"""Record MultiDiffusion runtime arguments."""
|
|
|
|
calls["multidiffusion"] = kwargs
|
|
return output
|
|
|
|
def fake_mixture(**kwargs: Any) -> dict[str, Any]:
|
|
"""Fail if Mixture runtime is selected."""
|
|
|
|
calls["mixture"] = kwargs
|
|
raise AssertionError("Mixture runtime should not be called.")
|
|
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"multidiffusion_sampling.sample_multidiffusion",
|
|
fake_multidiffusion,
|
|
)
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"mixture_of_diffusers_sampling.sample_mixture_of_diffusers",
|
|
fake_mixture,
|
|
)
|
|
|
|
result = TiledDiffusionSamplingService().sample(
|
|
**_sample_kwargs(diffusion_mode="multidiffusion")
|
|
)
|
|
|
|
assert result is output
|
|
assert "multidiffusion" in calls
|
|
assert "mixture" not in calls
|
|
|
|
|
|
def test_mixture_mode_routes_to_mixture_runtime(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Mixture of Diffusers mode calls only the Mixture runtime."""
|
|
|
|
calls: dict[str, dict[str, Any]] = {}
|
|
output: dict[str, Any] = {"samples": torch.ones((1, 4, 4, 4))}
|
|
|
|
def fake_multidiffusion(**kwargs: Any) -> dict[str, Any]:
|
|
"""Fail if MultiDiffusion runtime is selected."""
|
|
|
|
calls["multidiffusion"] = kwargs
|
|
raise AssertionError("MultiDiffusion runtime should not be called.")
|
|
|
|
def fake_mixture(**kwargs: Any) -> dict[str, Any]:
|
|
"""Record Mixture runtime arguments."""
|
|
|
|
calls["mixture"] = kwargs
|
|
return output
|
|
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"multidiffusion_sampling.sample_multidiffusion",
|
|
fake_multidiffusion,
|
|
)
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"mixture_of_diffusers_sampling.sample_mixture_of_diffusers",
|
|
fake_mixture,
|
|
)
|
|
|
|
result = TiledDiffusionSamplingService().sample(
|
|
**_sample_kwargs(diffusion_mode="mixture_of_diffusers")
|
|
)
|
|
|
|
assert result is output
|
|
assert "mixture" in calls
|
|
assert "multidiffusion" not in calls
|
|
|
|
|
|
def test_service_forwards_sampling_arguments_unchanged(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""The dispatcher preserves every sampler argument at the runtime boundary."""
|
|
|
|
calls: dict[str, Any] = {}
|
|
output: dict[str, Any] = {"samples": torch.ones((1, 4, 4, 4))}
|
|
preview_context = object()
|
|
|
|
def fake_multidiffusion(**kwargs: Any) -> dict[str, Any]:
|
|
"""Record forwarded arguments."""
|
|
|
|
calls.update(kwargs)
|
|
return output
|
|
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"multidiffusion_sampling.sample_multidiffusion",
|
|
fake_multidiffusion,
|
|
)
|
|
|
|
kwargs = _sample_kwargs(
|
|
diffusion_mode="multidiffusion",
|
|
preview_context=preview_context,
|
|
)
|
|
result = TiledDiffusionSamplingService().sample(**kwargs)
|
|
|
|
assert result is output
|
|
assert calls == {
|
|
key: value for key, value in kwargs.items() if key != "diffusion_mode"
|
|
}
|
|
|
|
|
|
def test_service_forwards_differential_diffusion_request(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""The dispatcher preserves differential-denoise-mask composition requests."""
|
|
|
|
calls: dict[str, Any] = {}
|
|
|
|
def fake_multidiffusion(**kwargs: Any) -> dict[str, Any]:
|
|
"""Record forwarded arguments."""
|
|
|
|
calls.update(kwargs)
|
|
return {"samples": torch.ones((1, 4, 4, 4))}
|
|
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"multidiffusion_sampling.sample_multidiffusion",
|
|
fake_multidiffusion,
|
|
)
|
|
|
|
TiledDiffusionSamplingService().sample(
|
|
**(
|
|
_sample_kwargs(diffusion_mode="multidiffusion")
|
|
| {"differential_diffusion": True}
|
|
)
|
|
)
|
|
|
|
assert calls["differential_diffusion"] is True
|
|
|
|
|
|
def test_invalid_mode_fails_before_runtime_call(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Unsupported modes are rejected before any runtime sampler is called."""
|
|
|
|
def fail_runtime(**kwargs: Any) -> dict[str, Any]:
|
|
"""Fail if validation does not stop dispatch."""
|
|
|
|
del kwargs
|
|
raise AssertionError("Runtime should not be called.")
|
|
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"multidiffusion_sampling.sample_multidiffusion",
|
|
fail_runtime,
|
|
)
|
|
monkeypatch.setattr(
|
|
"simple_syrup.services.tiled_diffusion_sampling_service."
|
|
"mixture_of_diffusers_sampling.sample_mixture_of_diffusers",
|
|
fail_runtime,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="diffusion_mode"):
|
|
TiledDiffusionSamplingService().sample(
|
|
**_sample_kwargs(diffusion_mode="full_latent")
|
|
)
|
|
|
|
|
|
def _sample_kwargs(
|
|
*,
|
|
diffusion_mode: str,
|
|
preview_context: object | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Return valid tiled diffusion sample arguments."""
|
|
|
|
return {
|
|
"diffusion_mode": diffusion_mode,
|
|
"model": "model",
|
|
"seed": 123,
|
|
"steps": 20,
|
|
"cfg": 7.0,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"positive": "positive",
|
|
"negative": "negative",
|
|
"latent_image": {"samples": torch.zeros((1, 4, 4, 4))},
|
|
"denoise": 0.8,
|
|
"latent_tile_width": 128,
|
|
"latent_tile_height": 80,
|
|
"latent_tile_overlap": 24,
|
|
"latent_tile_batch_size": 3,
|
|
"preview_context": preview_context,
|
|
"differential_diffusion": False,
|
|
}
|