refiners/tests/foundationals/latent_diffusion/test_multi_diffusion.py
Laurent f44ae150a7 improve/add MultiDiffusion and MultiUpscaler e2e tests
Co-authored-by: limiteinductive <benjamin@lagon.tech>
Co-authored-by: Cédric Deltheil <355031+deltheil@users.noreply.github.com>
2024-07-11 15:23:02 +02:00

35 lines
1.3 KiB
Python

import pytest
from refiners.foundationals.latent_diffusion.multi_diffusion import MultiDiffusion, Size
def test_generate_latent_tiles() -> None:
size = Size(height=128, width=128)
tile_size = Size(height=32, width=32)
tiles = MultiDiffusion.generate_latent_tiles(size=size, tile_size=tile_size)
assert len(tiles) == 25
tiles = MultiDiffusion.generate_latent_tiles(size=size, tile_size=tile_size, min_overlap=0)
assert len(tiles) == 16
size = Size(height=100, width=200)
tile_size = Size(height=32, width=32)
tiles = MultiDiffusion.generate_latent_tiles(size=size, tile_size=tile_size, min_overlap=2)
assert len(tiles) == 28
def test_generate_latent_tiles_small_size() -> None:
# Test when the size is smaller than the tile size
size = Size(height=32, width=32)
tile_size = Size(height=64, width=64)
tiles = MultiDiffusion.generate_latent_tiles(size=size, tile_size=tile_size)
assert len(tiles) == 1
assert Size(tiles[0].bottom - tiles[0].top, tiles[0].right - tiles[0].left) == size
def test_overlap_larger_tile_size() -> None:
with pytest.raises(AssertionError):
size = Size(height=128, width=128)
tile_size = Size(height=32, width=32)
MultiDiffusion.generate_latent_tiles(size=size, tile_size=tile_size, min_overlap=32)