373 lines
12 KiB
Python
373 lines
12 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
|
|
|
|
"""Characterize Regional Conditioning across sampling layouts and tile batches."""
|
|
|
|
from collections.abc import Callable
|
|
from typing import Any, cast
|
|
|
|
import comfy.conds
|
|
import comfy.sampler_helpers
|
|
import comfy.samplers
|
|
import pytest
|
|
import torch
|
|
|
|
from simple_syrup.domain.conditioning_batch import ConditioningBatch
|
|
from simple_syrup.domain.contextual_diffusion import (
|
|
ContextualDiffusionControls,
|
|
build_contextual_diffusion_plan,
|
|
)
|
|
from simple_syrup.domain.segs import BoundingBox, CropRegion, Segment
|
|
from simple_syrup.domain.tiled_diffusion import build_tiled_diffusion_plan
|
|
from simple_syrup.runtime.contextual_model_wrapper import (
|
|
ContextualDiffusionModelWrapper,
|
|
)
|
|
from simple_syrup.runtime.mixture_of_diffusers_sampling import (
|
|
MixtureOfDiffusersModelWrapper,
|
|
)
|
|
from simple_syrup.runtime.multidiffusion_sampling import MultiDiffusionModelWrapper
|
|
from simple_syrup.services.regional_conditioning_service import (
|
|
RegionalConditioningService,
|
|
)
|
|
|
|
ModelFunction = Callable[..., torch.Tensor]
|
|
ModelFunctionWrapper = Callable[[ModelFunction, dict[str, Any]], torch.Tensor]
|
|
|
|
|
|
class _FakePatcher:
|
|
"""Provide the hook lifecycle required by Comfy's conditioning evaluator."""
|
|
|
|
def prepare_hook_patches_current_keyframe(
|
|
self,
|
|
timestep: torch.Tensor,
|
|
hooks: object,
|
|
model_options: dict[str, Any],
|
|
) -> None:
|
|
"""Accept Comfy's current-keyframe preparation call."""
|
|
del timestep, hooks, model_options
|
|
|
|
def prepare_state(
|
|
self,
|
|
timestep: torch.Tensor,
|
|
model_options: dict[str, Any],
|
|
) -> None:
|
|
"""Accept Comfy's hook-state preparation call."""
|
|
del timestep, model_options
|
|
|
|
def get_free_memory(self, device: torch.device) -> float:
|
|
"""Report ample memory so compatible conditions may batch together."""
|
|
del device
|
|
return 1e12
|
|
|
|
def apply_hooks(
|
|
self,
|
|
*,
|
|
hooks: object | None,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Accept hook activation without introducing model patches."""
|
|
del hooks
|
|
return {}
|
|
|
|
|
|
class _RegionalPredictionModel:
|
|
"""Evaluate regional conditioning while recording runtime batch shapes."""
|
|
|
|
def __init__(self, wrapper: ModelFunctionWrapper | None = None) -> None:
|
|
"""Initialize the deterministic prediction model."""
|
|
self.current_patcher = _FakePatcher()
|
|
self.wrapper = wrapper
|
|
self.calls: list[tuple[int, ...]] = []
|
|
|
|
def memory_required(
|
|
self,
|
|
input_shape: tuple[int, ...] | list[int],
|
|
**requirements: Any,
|
|
) -> float:
|
|
"""Report a small deterministic memory requirement."""
|
|
del input_shape, requirements
|
|
return 1.0
|
|
|
|
def apply_model(
|
|
self,
|
|
input_x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
**conditioning: Any,
|
|
) -> torch.Tensor:
|
|
"""Apply the selected sampling wrapper around the deterministic model."""
|
|
if self.wrapper is None:
|
|
return self._evaluate(input_x, timestep, **conditioning)
|
|
return self.wrapper(
|
|
self._evaluate,
|
|
{"input": input_x, "timestep": timestep, "c": conditioning},
|
|
)
|
|
|
|
def _evaluate(
|
|
self,
|
|
input_x: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
**conditioning: Any,
|
|
) -> torch.Tensor:
|
|
"""Turn each cross-attention value into a spatially constant prediction."""
|
|
del timestep
|
|
self.calls.append(tuple(input_x.shape))
|
|
cross_attention = conditioning["c_crossattn"]
|
|
if not isinstance(cross_attention, torch.Tensor):
|
|
raise TypeError("Cross-attention conditioning must be a tensor.")
|
|
reduction_dimensions = tuple(range(1, cross_attention.ndim))
|
|
values = cross_attention.float().mean(dim=reduction_dimensions)
|
|
value_shape = (-1,) + (1,) * (input_x.ndim - 1)
|
|
return torch.ones_like(input_x) * values.reshape(value_shape)
|
|
|
|
|
|
def _conditioning(value: float) -> list[list[object]]:
|
|
"""Build one raw conditioning entry with a recognizable prediction value."""
|
|
cross_attention = torch.full((1, 2, 3), value, dtype=torch.float32)
|
|
return [[cross_attention, {}]]
|
|
|
|
|
|
def _regional_conditioning(mask: torch.Tensor) -> list[dict[str, Any]]:
|
|
"""Build converted global-plus-regional conditioning for Comfy evaluation."""
|
|
service = RegionalConditioningService()
|
|
positive = ConditioningBatch((_conditioning(1.0), _conditioning(3.0)))
|
|
negative = ConditioningBatch((_conditioning(-1.0), _conditioning(-3.0)))
|
|
regional_positive, _regional_negative = service.assemble(
|
|
positive=positive,
|
|
negative=negative,
|
|
masks=mask,
|
|
regional_prompt_weight=0.5,
|
|
region_mask_feather=0,
|
|
)
|
|
converted = cast(
|
|
list[dict[str, Any]], comfy.sampler_helpers.convert_cond(regional_positive)
|
|
)
|
|
for entry in converted:
|
|
cross_attention = entry.pop("cross_attn")
|
|
entry["model_conds"] = {
|
|
"c_crossattn": comfy.conds.CONDCrossAttn(cross_attention)
|
|
}
|
|
return converted
|
|
|
|
|
|
def _samples(layout_dimensions: int) -> torch.Tensor:
|
|
"""Create a BCHW or singleton-depth BCDHW latent."""
|
|
if layout_dimensions == 4:
|
|
return torch.zeros((1, 1, 32, 64), dtype=torch.float32)
|
|
return torch.zeros((1, 1, 1, 32, 64), dtype=torch.float32)
|
|
|
|
|
|
def _left_half_mask() -> torch.Tensor:
|
|
"""Create a mask selecting the left half of the latent plane."""
|
|
mask = torch.zeros((1, 32, 64), dtype=torch.float32)
|
|
mask[:, :, :32] = 1.0
|
|
return mask
|
|
|
|
|
|
def _expected_prediction(samples: torch.Tensor) -> torch.Tensor:
|
|
"""Create the blended prediction expected from the regional mask."""
|
|
expected = torch.ones_like(samples)
|
|
expected[..., :32] = 2.0
|
|
return expected
|
|
|
|
|
|
def _predict(
|
|
samples: torch.Tensor,
|
|
wrapper: ModelFunctionWrapper | None,
|
|
) -> tuple[torch.Tensor, _RegionalPredictionModel]:
|
|
"""Evaluate Regional Conditioning through Comfy's native condition path."""
|
|
return _predict_conditioning(
|
|
samples,
|
|
_regional_conditioning(_left_half_mask()),
|
|
wrapper,
|
|
)
|
|
|
|
|
|
def _predict_conditioning(
|
|
samples: torch.Tensor,
|
|
conditioning: list[dict[str, Any]],
|
|
wrapper: ModelFunctionWrapper | None,
|
|
) -> tuple[torch.Tensor, _RegionalPredictionModel]:
|
|
"""Evaluate one converted regional conditioning through an optional wrapper."""
|
|
|
|
model = _RegionalPredictionModel(wrapper)
|
|
prediction = comfy.samplers.calc_cond_batch(
|
|
model,
|
|
[conditioning],
|
|
samples,
|
|
torch.ones((samples.shape[0],), dtype=torch.float32),
|
|
{},
|
|
)[0]
|
|
return prediction, model
|
|
|
|
|
|
@pytest.mark.parametrize("layout_dimensions", [4, 5])
|
|
def test_full_sampling_preserves_regional_prediction(
|
|
layout_dimensions: int,
|
|
) -> None:
|
|
"""Preserve regional blending for full BCHW and BCDHW sampling."""
|
|
samples = _samples(layout_dimensions)
|
|
|
|
prediction, model = _predict(samples, None)
|
|
|
|
torch.testing.assert_close(prediction, _expected_prediction(samples))
|
|
assert prediction.shape == samples.shape
|
|
assert model.calls == [(2, *samples.shape[1:])]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"wrapper_type",
|
|
[MultiDiffusionModelWrapper, MixtureOfDiffusersModelWrapper],
|
|
)
|
|
@pytest.mark.parametrize("layout_dimensions", [4, 5])
|
|
@pytest.mark.parametrize("tile_batch_size", [1, 2, 4, 8])
|
|
def test_tiled_sampling_preserves_regional_prediction(
|
|
wrapper_type: type[MultiDiffusionModelWrapper]
|
|
| type[MixtureOfDiffusersModelWrapper],
|
|
layout_dimensions: int,
|
|
tile_batch_size: int,
|
|
) -> None:
|
|
"""Preserve regional blending across tiled modes, layouts, and batch sizes."""
|
|
plan = build_tiled_diffusion_plan(
|
|
latent_width=64,
|
|
latent_height=32,
|
|
tile_width=16,
|
|
tile_height=16,
|
|
overlap=0,
|
|
tile_batch_size=tile_batch_size,
|
|
)
|
|
wrapper = wrapper_type(plan=plan, existing_wrapper=None)
|
|
samples = _samples(layout_dimensions)
|
|
|
|
prediction, model = _predict(samples, wrapper)
|
|
|
|
torch.testing.assert_close(prediction, _expected_prediction(samples))
|
|
assert prediction.shape == samples.shape
|
|
assert max(call[0] for call in model.calls) == tile_batch_size * 2
|
|
|
|
|
|
@pytest.mark.parametrize("diffusion_mode", ["multidiffusion", "mixture_of_diffusers"])
|
|
@pytest.mark.parametrize("layout_dimensions", [4, 5])
|
|
@pytest.mark.parametrize("tile_batch_size", [1, 2, 4, 8])
|
|
def test_contextual_sampling_preserves_regional_prediction(
|
|
diffusion_mode: str,
|
|
layout_dimensions: int,
|
|
tile_batch_size: int,
|
|
) -> None:
|
|
"""Preserve regional blending across contextual modes and layouts."""
|
|
controls = ContextualDiffusionControls(
|
|
latent_context_size=16,
|
|
latent_context_overlap=0,
|
|
latent_context_batch_size=tile_batch_size,
|
|
global_weight=1.0,
|
|
global_steps=1,
|
|
global_decay=0.5,
|
|
)
|
|
plan = build_contextual_diffusion_plan(
|
|
latent_width=64,
|
|
latent_height=32,
|
|
controls=controls,
|
|
segs=None,
|
|
region_masks=None,
|
|
)
|
|
wrapper = ContextualDiffusionModelWrapper(
|
|
plan=plan,
|
|
controls=controls,
|
|
sigmas=torch.tensor([1.0, 0.0], dtype=torch.float32),
|
|
diffusion_mode=diffusion_mode,
|
|
existing_wrapper=None,
|
|
)
|
|
samples = _samples(layout_dimensions)
|
|
|
|
prediction, model = _predict(samples, wrapper)
|
|
|
|
torch.testing.assert_close(prediction, _expected_prediction(samples))
|
|
assert prediction.shape == samples.shape
|
|
assert max(call[0] for call in model.calls) == tile_batch_size * 2
|
|
|
|
|
|
@pytest.mark.parametrize("diffusion_mode", ["multidiffusion", "mixture_of_diffusers"])
|
|
@pytest.mark.parametrize("layout_dimensions", [4, 5])
|
|
@pytest.mark.parametrize("use_segs", [False, True], ids=["grid", "segs"])
|
|
def test_contextual_sampling_preserves_overlapping_and_uncovered_regions(
|
|
diffusion_mode: str,
|
|
layout_dimensions: int,
|
|
use_segs: bool,
|
|
) -> None:
|
|
"""Match native regional overlap and fallback across grid and SEGS plans."""
|
|
|
|
masks = _overlap_and_uncovered_masks()
|
|
conditioning = _regional_conditioning_for_masks(masks)
|
|
samples = _samples(layout_dimensions)
|
|
expected, _direct_model = _predict_conditioning(samples, conditioning, None)
|
|
controls = ContextualDiffusionControls(16, 4, 2, 1.0, 1, 0.5)
|
|
plan = build_contextual_diffusion_plan(
|
|
latent_width=64,
|
|
latent_height=32,
|
|
controls=controls,
|
|
segs=_context_segs() if use_segs else None,
|
|
region_masks=masks,
|
|
)
|
|
wrapper = ContextualDiffusionModelWrapper(
|
|
plan=plan,
|
|
controls=controls,
|
|
sigmas=torch.tensor([1.0, 0.0], dtype=torch.float32),
|
|
diffusion_mode=diffusion_mode,
|
|
existing_wrapper=None,
|
|
)
|
|
|
|
prediction, model = _predict_conditioning(samples, conditioning, wrapper)
|
|
|
|
torch.testing.assert_close(prediction, expected)
|
|
assert prediction.shape == samples.shape
|
|
assert all(tile.weight_mask is not None for tile in plan.tile_plan.tiles)
|
|
assert len(model.calls) > len(_direct_model.calls)
|
|
|
|
|
|
def _overlap_and_uncovered_masks() -> torch.Tensor:
|
|
"""Create two overlapping regions with uncovered strips on both sides."""
|
|
|
|
masks = torch.zeros((2, 32, 64), dtype=torch.float32)
|
|
masks[0, :, 8:36] = 1.0
|
|
masks[1, :, 28:56] = 1.0
|
|
return masks
|
|
|
|
|
|
def _regional_conditioning_for_masks(
|
|
masks: torch.Tensor,
|
|
) -> list[dict[str, Any]]:
|
|
"""Build converted global, first-region, and second-region conditioning."""
|
|
|
|
positive, _negative = RegionalConditioningService().assemble(
|
|
positive=ConditioningBatch(
|
|
(_conditioning(1.0), _conditioning(3.0), _conditioning(5.0))
|
|
),
|
|
negative=_conditioning(-1.0),
|
|
masks=masks,
|
|
regional_prompt_weight=0.5,
|
|
region_mask_feather=0,
|
|
)
|
|
converted = cast(list[dict[str, Any]], comfy.sampler_helpers.convert_cond(positive))
|
|
for entry in converted:
|
|
cross_attention = entry.pop("cross_attn")
|
|
entry["model_conds"] = {
|
|
"c_crossattn": comfy.conds.CONDCrossAttn(cross_attention)
|
|
}
|
|
return converted
|
|
|
|
|
|
def _context_segs() -> tuple[tuple[int, int], tuple[Segment, ...]]:
|
|
"""Create one central SEGS constraint over the regional planning canvas."""
|
|
|
|
mask = torch.zeros((32, 64), dtype=torch.float32)
|
|
mask[4:28, 12:52] = 1.0
|
|
crop = CropRegion(0, 0, 64, 32)
|
|
segment = Segment(
|
|
cropped_image=None,
|
|
cropped_mask=mask,
|
|
confidence=1.0,
|
|
crop_region=crop,
|
|
bbox=BoundingBox(*crop),
|
|
label="context",
|
|
)
|
|
return (32, 64), (segment,)
|