feat: #678 opts

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()

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