diff --git a/adetailer/opts.py b/adetailer/opts.py new file mode 100644 index 0000000..e7f0c59 --- /dev/null +++ b/adetailer/opts.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import ClassVar, TypeVar + +import numpy as np + +T = TypeVar("T", int, float) + + +def get_dynamic_denoise_strength( + denoise_power: float, + denoise_strength: float, + bbox: Sequence[T], + image_size: tuple[int, int], +) -> float: + if np.isclose(denoise_power, 0.0) or len(bbox) != 4: + return denoise_strength + + width, height = image_size + + image_pixels = width * height + bbox_pixels = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) + + normalized_area = bbox_pixels / image_pixels + denoise_modifier = (1.0 - normalized_area) ** denoise_power + + return denoise_strength * denoise_modifier + + +class _OptimalCropSize: + sdxl_res: ClassVar[list[tuple[int, int]]] = [ + (1024, 1024), + (1152, 896), + (896, 1152), + (1216, 832), + (832, 1216), + (1344, 768), + (768, 1344), + (1536, 640), + (640, 1536), + ] + + def sdxl( + self, inpaint_width: int, inpaint_height: int, bbox: Sequence[T] + ) -> tuple[int, int]: + if len(bbox) != 4: + return inpaint_width, inpaint_height + + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + bbox_aspect_ratio = bbox_width / bbox_height + + resolutions = [ + res + for res in self.sdxl_res + if (res[0] >= bbox_width and res[1] >= bbox_height) + and (res[0] >= inpaint_width or res[1] >= inpaint_height) + ] + + if not resolutions: + return inpaint_width, inpaint_height + + return min( + resolutions, + key=lambda res: abs((res[0] / res[1]) - bbox_aspect_ratio), + ) + + def free( + self, inpaint_width: int, inpaint_height: int, bbox: Sequence[T] + ) -> tuple[int, int]: + if len(bbox) != 4: + return inpaint_width, inpaint_height + + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + bbox_aspect_ratio = bbox_width / bbox_height + + scale_size = max(inpaint_width, inpaint_height) + + if bbox_aspect_ratio > 1: + optimal_width = scale_size + optimal_height = scale_size / bbox_aspect_ratio + else: + optimal_width = scale_size * bbox_aspect_ratio + optimal_height = scale_size + + # Round up to the nearest multiple of 8 to make the dimensions friendly for upscaling/diffusion. + optimal_width = ((optimal_width + 8 - 1) // 8) * 8 + optimal_height = ((optimal_height + 8 - 1) // 8) * 8 + + return int(optimal_width), int(optimal_height) + + +optimal_crop_size = _OptimalCropSize() diff --git a/tests/test_opts.py b/tests/test_opts.py new file mode 100644 index 0000000..347e3b9 --- /dev/null +++ b/tests/test_opts.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import numpy as np +import pytest +from hypothesis import assume, given +from hypothesis import strategies as st + +from adetailer.opts import get_dynamic_denoise_strength, optimal_crop_size + + +@pytest.mark.parametrize( + ("denoise_power", "denoise_strength", "bbox", "image_size", "expected_result"), + [ + (0.001, 0.5, [0, 0, 100, 100], (200, 200), 0.4998561796520339), + (1.5, 0.3, [0, 0, 100, 100], (200, 200), 0.1948557158514987), + (-0.001, 0.7, [0, 0, 100, 100], (1000, 1000), 0.7000070352704507), + (-0.5, 0.5, [0, 0, 100, 100], (1000, 1000), 0.502518907629606), + ], +) +def test_get_dynamic_denoise_strength( + denoise_power: float, + denoise_strength: float, + bbox: list[int], + image_size: tuple[int, int], + expected_result: float, +): + result = get_dynamic_denoise_strength( + denoise_power, denoise_strength, bbox, image_size + ) + assert np.isclose(result, expected_result) + + +@given(denoise_strength=st.floats(allow_nan=False)) +def test_get_dynamic_denoise_strength_no_bbox(denoise_strength: float): + result = get_dynamic_denoise_strength(0.5, denoise_strength, [], (1000, 1000)) + assert result == denoise_strength + + +@given(denoise_strength=st.floats(allow_nan=False)) +def test_get_dynamic_denoise_strength_zero_power(denoise_strength: float): + result = get_dynamic_denoise_strength(0.0, denoise_strength, [], (1000, 1000)) + assert result == denoise_strength + + +@given( + inpaint_width=st.integers(1), + inpaint_height=st.integers(1), + bbox=st.tuples( + st.integers(0, 500), + st.integers(0, 500), + st.integers(501, 1000), + st.integers(501, 1000), + ), +) +def test_optimal_crop_size_sdxl( + inpaint_width: int, inpaint_height: int, bbox: tuple[int, int, int, int] +): + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + assume(bbox_width > 0 and bbox_height > 0) + + result = optimal_crop_size.sdxl(inpaint_width, inpaint_height, bbox) + assert (result in optimal_crop_size.sdxl_res) or result == ( + inpaint_width, + inpaint_height, + ) + + if result != (inpaint_width, inpaint_height): + assert result[0] >= bbox_width + assert result[1] >= bbox_height + assert result[0] >= inpaint_width or result[1] >= inpaint_height + + +@given( + inpaint_width=st.integers(1), + inpaint_height=st.integers(1), + bbox=st.tuples( + st.integers(0, 500), + st.integers(0, 500), + st.integers(501, 1000), + st.integers(501, 1000), + ), +) +def test_optimal_crop_size_free( + inpaint_width: int, inpaint_height: int, bbox: tuple[int, int, int, int] +): + bbox_width = bbox[2] - bbox[0] + bbox_height = bbox[3] - bbox[1] + assume(bbox_width > 0 and bbox_height > 0) + + result = optimal_crop_size.free(inpaint_width, inpaint_height, bbox) + assert result[0] % 8 == 0 + assert result[1] % 8 == 0