# SimpleSyrup - workflow-focused ComfyUI extensions for image generation # Copyright (C) 2026 Artificial Sweetener and contributors # SPDX-License-Identifier: AGPL-3.0-or-later """Verify exact generic regional convolution execution against references.""" from __future__ import annotations import pytest import torch import torch.nn.functional as functional from comfy.weight_adapter.lora import LoRAAdapter from torch import nn from simple_syrup.domain.regional_activation_geometry import ( RegionalActivationBatchAlignment, RegionalActivationGeometry, RegionalActivationLayout, RegionalTemporalOwnership, ) from simple_syrup.domain.spatial_views import ( SpatialBatchLayout, SpatialView, SpatialViewKind, ) from simple_syrup.masking.regional_activation_mask_projection import ( RegionalActivationMaskBatch, ) from simple_syrup.runtime.regional_lora.convolution_execution import ( RegionalConvolutionExecutor, ) from simple_syrup.runtime.regional_lora.convolution_execution_plan import ( RegionalConvolutionExecutionPlan, RegionalConvolutionParameters, RegionalConvolutionTargetUse, ) from simple_syrup.runtime.regional_lora.convolution_preparation import ( RegionalConvolutionPreparation, ) @pytest.mark.parametrize("dimension", [1, 2, 3]) def test_direct_convolution_dimensions_match_rank_space_reference( dimension: int, ) -> None: """Match exact Conv1d/2d/3d down-mask-up math with spatial kernels.""" fixture = _fixture(dimension=dimension, kernel_size=3, padding=1) mask = torch.linspace(0.0, 1.0, fixture.rank_element_count).reshape( 1, *fixture.rank_spatial, ) result = fixture.execute(mask, schedule=0.6) rank = _conv( fixture.inputs, fixture.down, stride=fixture.parameters.stride, padding=fixture.parameters.padding, dilation=fixture.parameters.dilation, groups=fixture.parameters.groups, ) expected = fixture.original(fixture.inputs) + _conv( rank * mask.unsqueeze(1) * 0.8 * 0.6, fixture.up, stride=(1,) * dimension, padding=(0,) * dimension, dilation=(1,) * dimension, groups=1, ) torch.testing.assert_close(result, expected, rtol=1e-6, atol=1e-6) def test_pointwise_convolution_all_one_matches_global_delta() -> None: """Match ordinary complete-image convolution LoRA at full coverage.""" fixture = _fixture(dimension=2, kernel_size=1, padding=0) result = fixture.execute(torch.ones((1, *fixture.rank_spatial))) expected = fixture.original(fixture.inputs) + _conv( _conv( fixture.inputs, fixture.down, stride=fixture.parameters.stride, padding=fixture.parameters.padding, dilation=fixture.parameters.dilation, groups=1, ) * 0.8, fixture.up, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) torch.testing.assert_close(result, expected, rtol=1e-6, atol=1e-6) def test_locon_middle_applies_mask_after_spatial_middle() -> None: """Project pointwise down through middle kernel before regional masking.""" fixture = _fixture(dimension=2, kernel_size=3, padding=1, use_middle=True) mask = torch.linspace(0.2, 0.9, fixture.rank_element_count).reshape( 1, *fixture.rank_spatial, ) result = fixture.execute(mask) pointwise = _conv( fixture.inputs, fixture.down, stride=fixture.parameters.stride, padding=fixture.parameters.padding, dilation=fixture.parameters.dilation, groups=1, ) assert fixture.middle is not None rank = _conv( pointwise, fixture.middle, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) expected = fixture.original(fixture.inputs) + _conv( rank * mask.unsqueeze(1) * 0.8, fixture.up, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) torch.testing.assert_close(result, expected, rtol=1e-6, atol=1e-6) def test_stride_changes_rank_mask_and_output_resolution_exactly() -> None: """Require masks at down-convolution output resolution, not input resolution.""" fixture = _fixture(dimension=2, kernel_size=3, padding=1, stride=2) assert fixture.inputs.shape[-2:] != fixture.rank_spatial result = fixture.execute(torch.ones((1, *fixture.rank_spatial))) assert tuple(result.shape[-2:]) == fixture.rank_spatial with pytest.raises(ValueError, match="exact rank activation shape"): fixture.execute(torch.ones((1, *fixture.inputs.shape[-2:]))) def test_zero_mask_and_schedule_skip_preparation_but_call_original_once() -> None: """Avoid device preparation for inactive convolution adapters.""" fixture = _fixture(dimension=2, kernel_size=3, padding=1) original = _CountedOriginal(fixture.original) result = fixture.execute( torch.zeros((1, *fixture.rank_spatial)), original=original, ) assert torch.equal(result, fixture.original(fixture.inputs)) assert original.calls == 1 assert fixture.use.preparation._prepared == {} def test_multiple_adapters_preserve_declared_addition_order() -> None: """Add distinct convolution targets in declared schedule order.""" fixture = _fixture(dimension=2, kernel_size=3, padding=1) second_down = fixture.down * -0.5 second_up = fixture.up * 0.25 second = RegionalConvolutionTargetUse( 1, 1, (id(second_down), id(second_up), None), RegionalConvolutionPreparation(second_down, None, second_up), fixture.parameters, 0.4, ) plan = RegionalConvolutionExecutionPlan((fixture.use, second)) masks = torch.stack( ( torch.full(fixture.rank_spatial, 0.25), torch.full(fixture.rank_spatial, 0.75), ) ) result = fixture.execute(masks, schedules=(0.5, 0.75), plan=plan) expected = fixture.original(fixture.inputs) for down, up, mask, scale in ( (fixture.down, fixture.up, masks[0], 0.8 * 0.5), (second_down, second_up, masks[1], 0.4 * 0.75), ): rank = _conv( fixture.inputs, down, stride=fixture.parameters.stride, padding=fixture.parameters.padding, dilation=fixture.parameters.dilation, groups=1, ) expected = expected + _conv( rank * mask.unsqueeze(0).unsqueeze(0) * scale, up, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) torch.testing.assert_close(result, expected, rtol=1e-6, atol=1e-6) def test_grouped_convolution_matches_grouped_merged_weight_reference() -> None: """Repeat rank bases per group so all-one execution equals merged weight math.""" original = nn.Conv2d(4, 6, 3, padding=1, groups=2, bias=True) inputs = torch.linspace(-1.0, 1.0, 2 * 4 * 5 * 5).reshape(2, 4, 5, 5) down = torch.linspace(-0.2, 0.3, 2 * 2 * 3 * 3).reshape(2, 2, 3, 3) up = torch.linspace(-0.4, 0.25, 6 * 2).reshape(6, 2, 1, 1) parameters = RegionalConvolutionParameters( 2, (1, 1), (1, 1), (1, 1), 2, 4, 6, (3, 3), ) use = RegionalConvolutionTargetUse( 0, 0, (id(down), id(up), None), RegionalConvolutionPreparation(down, None, up), parameters, 0.75, ) plan = RegionalConvolutionExecutionPlan((use,)) masks = _mask_batch(torch.ones((1, 5, 5)), dimension=2, input_batch=2, features=4) result = RegionalConvolutionExecutor().execute( original, inputs, plan=plan, masks=masks, schedule_strengths=(1.0,), ) adapter = LoRAAdapter( {"up", "down"}, (up, down, None, None, None, None), ) patched = adapter.calculate_weight( original.weight.detach().clone(), "diffusion_model.grouped.weight", 0.75, 1.0, None, lambda value: value, intermediate_dtype=torch.float32, ) expected = functional.conv2d( inputs, patched, original.bias, stride=1, padding=1, groups=2, ) torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-6) def test_view_major_tiled_batch_keeps_each_tile_mask_on_its_rows() -> None: """Apply distinct view masks over existing view/chunk/latent batch order.""" fixture = _fixture(dimension=2, kernel_size=1, padding=0) fixture.inputs = fixture.inputs.repeat(2, 1, 1, 1) fixture.rank_spatial = tuple(int(value) for value in fixture.inputs.shape[2:]) view_values = torch.cat( ( torch.full((2, *fixture.rank_spatial), 0.25), torch.full((2, *fixture.rank_spatial), 0.75), ), dim=0, ) masks = RegionalActivationMaskBatch( view_values.unsqueeze(0).unsqueeze(2), RegionalActivationGeometry( RegionalActivationLayout.DIRECT_CONVOLUTION_2D, (4, 2, *fixture.rank_spatial), 1, fixture.rank_spatial[0], fixture.rank_spatial[1], RegionalActivationBatchAlignment(2, 1, _two_view_layout()), ), ) result = RegionalConvolutionExecutor().execute( fixture.original, fixture.inputs, plan=fixture.plan, masks=masks, schedule_strengths=(1.0,), ) rank = _conv( fixture.inputs, fixture.down, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) expected = fixture.original(fixture.inputs) + _conv( rank * view_values.unsqueeze(1) * 0.8, fixture.up, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) torch.testing.assert_close(result, expected, rtol=1e-6, atol=1e-6) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable") @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) def test_cuda_low_precision_matches_explicit_reference(dtype: torch.dtype) -> None: """Preserve full-rank convolution LoRA quality in FP16 and BF16.""" fixture = _fixture( dimension=2, kernel_size=3, padding=1, device=torch.device("cuda"), dtype=dtype, ) mask = torch.linspace( 0.0, 1.0, fixture.rank_element_count, device="cuda", dtype=dtype, ).reshape(1, *fixture.rank_spatial) result = fixture.execute(mask) rank = _conv( fixture.inputs, fixture.down.to("cuda", dtype), stride=fixture.parameters.stride, padding=fixture.parameters.padding, dilation=fixture.parameters.dilation, groups=1, ) expected = fixture.original(fixture.inputs) + _conv( rank * mask.unsqueeze(1) * 0.8, fixture.up.to("cuda", dtype), stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, ) torch.testing.assert_close(result, expected, rtol=0.02, atol=0.02) class _Fixture: """Retain one deterministic dimensional convolution execution case.""" def __init__( self, *, dimension: int, kernel_size: int, padding: int, stride: int, use_middle: bool, device: torch.device, dtype: torch.dtype, ) -> None: """Create exact original and low-rank convolution tensors.""" convolution = (nn.Conv1d, nn.Conv2d, nn.Conv3d)[dimension - 1] spatial = (7,) * dimension self.original = convolution( 2, 3, kernel_size, stride=stride, padding=padding, bias=True, device=device, dtype=dtype, ) self.inputs = torch.linspace( -1.0, 1.0, 2 * 2 * (7**dimension), device=device, dtype=dtype, ).reshape(2, 2, *spatial) rank = 2 down_kernel = (1,) * dimension if use_middle else (kernel_size,) * dimension self.down = torch.linspace( -0.3, 0.4, rank * 2 * (down_kernel[0] ** dimension), ).reshape(rank, 2, *down_kernel) self.middle = ( torch.linspace(-0.2, 0.25, rank * rank * (kernel_size**dimension)).reshape( rank, rank, *((kernel_size,) * dimension), ) if use_middle else None ) self.up = torch.linspace(-0.4, 0.35, 3 * rank).reshape( 3, rank, *((1,) * dimension), ) self.parameters = RegionalConvolutionParameters( dimension, (stride,) * dimension, (padding,) * dimension, (1,) * dimension, 1, 2, 3, (kernel_size,) * dimension, ) self.use = RegionalConvolutionTargetUse( 0, 0, ( id(self.down), id(self.up), None if self.middle is None else id(self.middle), ), RegionalConvolutionPreparation(self.down, self.middle, self.up), self.parameters, 0.8, ) self.plan = RegionalConvolutionExecutionPlan((self.use,)) self.rank_spatial = _reference_rank_spatial(self.inputs, self.use) self.rank_element_count = int(torch.tensor(self.rank_spatial).prod()) def execute( self, masks: torch.Tensor, *, schedule: float = 1.0, schedules: tuple[float, ...] | None = None, plan: RegionalConvolutionExecutionPlan | None = None, original: object | None = None, ) -> torch.Tensor: """Execute this fixture with explicitly shaped rank masks.""" selected_plan = plan or self.plan mask_batch = _mask_batch( masks.to(self.inputs), dimension=self.parameters.dimension, input_batch=int(self.inputs.shape[0]), features=2, ) selected_original = self.original if original is None else original assert callable(selected_original) return RegionalConvolutionExecutor().execute( selected_original, self.inputs, plan=selected_plan, masks=mask_batch, schedule_strengths=schedules or (schedule,), ) class _CountedOriginal: """Count original convolution calls while preserving output behavior.""" def __init__(self, original: nn.Module) -> None: """Retain the original module and zero calls.""" self._original = original self.calls = 0 def __call__(self, inputs: torch.Tensor) -> torch.Tensor: """Count and forward one exact original call.""" self.calls += 1 result = self._original(inputs) assert isinstance(result, torch.Tensor) return result def _fixture( *, dimension: int, kernel_size: int, padding: int, stride: int = 1, use_middle: bool = False, device: torch.device | None = None, dtype: torch.dtype = torch.float32, ) -> _Fixture: """Return one deterministic convolution execution fixture.""" return _Fixture( dimension=dimension, kernel_size=kernel_size, padding=padding, stride=stride, use_middle=use_middle, device=device or torch.device("cpu"), dtype=dtype, ) def _mask_batch( values: torch.Tensor, *, dimension: int, input_batch: int, features: int, ) -> RegionalActivationMaskBatch: """Wrap region/spatial masks in direct convolution geometry.""" if values.ndim == dimension: values = values.unsqueeze(0) regions = int(values.shape[0]) spatial = tuple(int(value) for value in values.shape[1:]) invocation = (input_batch, features, *spatial) layout = { 1: RegionalActivationLayout.DIRECT_CONVOLUTION_1D, 2: RegionalActivationLayout.DIRECT_CONVOLUTION_2D, 3: RegionalActivationLayout.DIRECT_CONVOLUTION_3D, }[dimension] geometry = RegionalActivationGeometry( layout, invocation, 1, 1 if dimension == 1 else spatial[-2], spatial[-1], RegionalActivationBatchAlignment(input_batch, 1), temporal_axis=2 if dimension == 3 else None, temporal_ownership=( RegionalTemporalOwnership.REPEAT_SPATIAL_MASK if dimension == 3 else RegionalTemporalOwnership.NONE ), ) expanded = values.reshape(regions, 1, 1, *spatial).expand( -1, input_batch, -1, *spatial, ) return RegionalActivationMaskBatch(expanded, geometry) def _reference_rank_spatial( inputs: torch.Tensor, use: RegionalConvolutionTargetUse, ) -> tuple[int, ...]: """Observe the exact explicit down/middle rank activation shape.""" rank = _conv( inputs, use.preparation.down.to(inputs), stride=use.parameters.stride, padding=use.parameters.padding, dilation=use.parameters.dilation, groups=use.parameters.groups, ) if use.preparation.middle is not None: rank = _conv( rank, use.preparation.middle.to(inputs), stride=(1,) * use.parameters.dimension, padding=(0,) * use.parameters.dimension, dilation=(1,) * use.parameters.dimension, groups=1, ) return tuple(int(value) for value in rank.shape[2:]) def _conv( inputs: torch.Tensor, weight: torch.Tensor, *, stride: tuple[int, ...], padding: tuple[int, ...], dilation: tuple[int, ...], groups: int, ) -> torch.Tensor: """Dispatch one explicit dimensional reference convolution.""" function = (functional.conv1d, functional.conv2d, functional.conv3d)[ inputs.ndim - 3 ] return function(inputs, weight, None, stride, padding, dilation, groups) def _two_view_layout() -> SpatialBatchLayout: """Return two equal view-major tiles over a four-entry active batch.""" return SpatialBatchLayout( 14, 7, ( SpatialView(SpatialViewKind.TILE, 0, 0, 7, 7, 7, 7), SpatialView(SpatialViewKind.TILE, 7, 0, 7, 7, 7, 7), ), input_batch_size=2, )