Artificial-Sweetener-SimpleSyrup/tests/test_segs_tiled_diffusion.py

# 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 SEGS-guided irregular tiled diffusion planning."""

from __future__ import annotations

import pytest
import torch

from simple_syrup.domain.segs import BoundingBox, CropRegion, Segment
from simple_syrup.domain.segs_tiled_diffusion import (
    _segment_mask_to_latent,
    build_segs_guided_tiled_diffusion_plan,
)


def test_guided_plan_splits_oversized_region_with_bounded_rectangular_windows() -> None:
    """A large region becomes multiple bounded cores instead of one giant tile."""

    plan = build_segs_guided_tiled_diffusion_plan(
        segs=_full_mask_segs(64, 64),
        latent_width=64,
        latent_height=64,
        tile_width=32,
        tile_height=32,
        overlap=8,
        tile_batch_size=3,
    )

    assert len(plan.tiles) > 1
    assert all(tile.width == 32 and tile.height == 32 for tile in plan.tiles)
    assert all(tile.weight_mask is not None for tile in plan.tiles)
    coverage = torch.zeros((64, 64), dtype=torch.float32)
    for tile in plan.tiles:
        assert tile.weight_mask is not None
        coverage[tile.y : tile.y + tile.height, tile.x : tile.x + tile.width] += (
            tile.weight_mask
        )
    assert bool((coverage > 0).all())


def test_guided_plan_feathers_the_configured_overlap_between_irregular_cores() -> None:
    """Non-zero overlap produces fractional shared ownership weights."""

    plan = build_segs_guided_tiled_diffusion_plan(
        segs=_half_mask_segs(64, 64),
        latent_width=64,
        latent_height=64,
        tile_width=32,
        tile_height=32,
        overlap=8,
        tile_batch_size=4,
    )

    weights = [tile.weight_mask for tile in plan.tiles if tile.weight_mask is not None]
    assert any(bool(torch.any((weight > 0) & (weight < 1))) for weight in weights)


def test_guided_plan_prefers_smaller_overlapping_segs_for_ownership() -> None:
    """A nested small SEG remains an ownership core instead of being swallowed."""

    full = torch.ones((32, 32), dtype=torch.float32)
    small = torch.zeros((32, 32), dtype=torch.float32)
    small[8:24, 8:24] = 1.0
    segs = (
        (32, 32),
        (
            _segment(full, "large", 0.6),
            _segment(small, "small", 0.9),
        ),
    )

    plan = build_segs_guided_tiled_diffusion_plan(
        segs=segs,
        latent_width=32,
        latent_height=32,
        tile_width=32,
        tile_height=32,
        overlap=0,
        tile_batch_size=4,
    )

    assert len(plan.tiles) == 2
    assert all(tile.weight_mask is not None for tile in plan.tiles)


def test_guided_plan_merges_small_cores_without_comparing_tensor_values() -> None:
    """Small nearby cores merge through their indexes instead of tensor equality."""

    first = torch.zeros((32, 32), dtype=torch.float32)
    second = torch.zeros((32, 32), dtype=torch.float32)
    first[4:8, 4:8] = 1.0
    second[8:12, 8:12] = 1.0
    segs = (
        (32, 32),
        (
            _segment(first, "first", 1.0),
            _segment(second, "second", 1.0),
        ),
    )

    plan = build_segs_guided_tiled_diffusion_plan(
        segs=segs,
        latent_width=32,
        latent_height=32,
        tile_width=32,
        tile_height=32,
        overlap=0,
        tile_batch_size=4,
    )

    assert len(plan.tiles) == 1


def test_crop_local_mask_projects_directly_to_latent_space() -> None:
    """A small crop maps without materializing a full source-resolution mask."""

    crop = CropRegion(512, 256, 768, 512)
    segment = Segment(
        cropped_image=None,
        cropped_mask=torch.ones((256, 256), dtype=torch.float32),
        confidence=1.0,
        crop_region=crop,
        bbox=BoundingBox(*crop),
        label="small_region",
    )

    latent_mask = _segment_mask_to_latent(
        segment,
        source_height=4096,
        source_width=4096,
        latent_width=512,
        latent_height=512,
    )

    assert int(latent_mask.sum().item()) == 32 * 32
    assert bool(latent_mask[32:64, 64:96].all())
    assert not bool(latent_mask[:32].any())
    assert not bool(latent_mask[:, :64].any())


def test_guided_plan_rejects_mismatched_image_aspect_ratio() -> None:
    """SEGS from a different image fail before tiled sampling begins."""

    with pytest.raises(ValueError, match="aspect ratio"):
        build_segs_guided_tiled_diffusion_plan(
            segs=_full_mask_segs(16, 32),
            latent_width=32,
            latent_height=32,
            tile_width=16,
            tile_height=16,
            overlap=4,
            tile_batch_size=2,
        )


def _full_mask_segs(
    height: int, width: int
) -> tuple[tuple[int, int], tuple[Segment, ...]]:
    """Return one full-image SEG payload."""

    return ((height, width), (_segment(torch.ones((height, width)), "region", 1.0),))


def _half_mask_segs(
    height: int, width: int
) -> tuple[tuple[int, int], tuple[Segment, ...]]:
    """Return one left-half SEG payload with implicit background ownership."""

    mask = torch.zeros((height, width), dtype=torch.float32)
    mask[:, : width // 2] = 1.0
    return ((height, width), (_segment(mask, "left", 1.0),))


def _segment(mask: torch.Tensor, label: str, confidence: float) -> Segment:
    """Build one full-image Segment whose crop equals the source dimensions."""

    height, width = mask.shape
    region = CropRegion(0, 0, width, height)
    return Segment(
        cropped_image=None,
        cropped_mask=mask,
        confidence=confidence,
        crop_region=region,
        bbox=BoundingBox(0, 0, width, height),
        label=label,
    )