diff --git a/.gitignore b/.gitignore index 0b09e1e..689264b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,2 @@ .idea -.pytest_cache \ No newline at end of file +.pytest_cache diff --git a/README.md b/README.md index 524a73a..b706c14 100644 --- a/README.md +++ b/README.md @@ -27,3 +27,10 @@ Install via the extensions tab on the [AUTOMATIC1111 webui](https://github.com/A - i.e `25, 50, 75, 125, 150, 175, 200` `50, 125, 300` etc. ![settings.png](docs%2Fsettings.png) + +## Contributing: + +- Open to suggestions +- Pull requests are appreciated +- Write tests if possible and useful +- Run pre-commit! diff --git a/aspect_ratio_helper/_util.py b/aspect_ratio_helper/_util.py new file mode 100644 index 0000000..36d6701 --- /dev/null +++ b/aspect_ratio_helper/_util.py @@ -0,0 +1,57 @@ +def _display_multiplication(num): + return f'x{round(num / 100, 3)}' + + +def _display_raw_percentage(num): + return f'{num}%' + + +def _display_minus_and_plus(num): + num -= 100 + if num > 0: + return f'+{num}%' + return f'{num}%' + + +_DEFAULT_DISPLAY_KEY = 'Incremental/decremental percentage (-50%, +50%)' +_PREDEFINED_PERCENTAGES_DISPLAY_MAP = { + _DEFAULT_DISPLAY_KEY: _display_minus_and_plus, + 'Raw percentage (50%, 150%)': _display_raw_percentage, + 'Multiplication (x0.5, x1.5)': _display_multiplication, +} + +_MIN_DIMENSION = 64 +_MAX_DIMENSION = 2048 + + +def _scale_by_percentage(width, height, pct): + aspect_ratio = float(width) / float(height) + step = (pct - 1.0) + new_width = int(round(width * (1.0 + step))) + new_height = int(round(new_width / aspect_ratio)) + if new_width > _MAX_DIMENSION: + new_width = _MAX_DIMENSION + new_height = int(round(new_width / aspect_ratio)) + if new_height > _MAX_DIMENSION: + new_height = _MAX_DIMENSION + new_width = int(round(new_height * aspect_ratio)) + if new_width < _MIN_DIMENSION: + new_width = _MIN_DIMENSION + new_height = int(round(new_width / aspect_ratio)) + if new_height < _MIN_DIMENSION: + new_height = _MIN_DIMENSION + new_width = int(round(new_height * aspect_ratio)) + return new_width, new_height + + +def _scale_dimensions_to_max_dimension(width, height, max_dim): + if max_dim < _MIN_DIMENSION: + max_dim = _MIN_DIMENSION + elif max_dim > _MAX_DIMENSION: + max_dim = _MAX_DIMENSION + if max_dim == max(width, height): + return width, height + aspect_ratio = float(width) / float(height) + if width > height: + return max_dim, max(int(round(max_dim / aspect_ratio)), 1) + return max(int(round(max_dim * aspect_ratio)), 1), max_dim diff --git a/aspect_ratio_helper/main.py b/aspect_ratio_helper/main.py new file mode 100644 index 0000000..31b01b1 --- /dev/null +++ b/aspect_ratio_helper/main.py @@ -0,0 +1,155 @@ +import contextlib +from functools import partial + +import gradio as gr +from modules import script_callbacks +from modules import scripts +from modules import shared +from modules.shared import opts + +from aspect_ratio_helper._util import _DEFAULT_DISPLAY_KEY +from aspect_ratio_helper._util import _MAX_DIMENSION +from aspect_ratio_helper._util import _MIN_DIMENSION +from aspect_ratio_helper._util import _PREDEFINED_PERCENTAGES_DISPLAY_MAP +from aspect_ratio_helper._util import _scale_by_percentage +from aspect_ratio_helper._util import _scale_dimensions_to_max_dimension + + +_EXTENSION_NAME = 'Aspect Ratio Helper' + + +def on_ui_settings(): + section = 'aspect_ratio_helper', _EXTENSION_NAME + shared.opts.add_option( + key='arh_expand_by_default', + info=shared.OptionInfo( + default=False, + label='Expand by default', + section=section, + ), + ) + shared.opts.add_option( + key='arh_show_max_width_or_height', + info=shared.OptionInfo( + default=True, + label='Show maximum width or height button', + section=section, + ), + ) + shared.opts.add_option( + key='arh_max_width_or_height', + info=shared.OptionInfo( + default=_MAX_DIMENSION / 2, + label='Maximum width or height default', + component=gr.Slider, + component_args={ + 'minimum': _MIN_DIMENSION, + 'maximum': _MAX_DIMENSION, + 'step': 1, + }, + section=section, + ), + ) + shared.opts.add_option( + key='arh_show_predefined_percentages', + info=shared.OptionInfo( + default=True, + label='Show predefined percentage buttons', + section=section, + ), + ) + shared.opts.add_option( + key='arh_predefined_percentages', + info=shared.OptionInfo( + default='25, 50, 75, 125, 150, 175, 200', + label='Predefined percentage buttons, applied to dimensions (75, ' + '125, 150)', + section=section, + ), + ) + shared.opts.add_option( + key='arh_predefined_percentages_display_key', + info=shared.OptionInfo( + default=_DEFAULT_DISPLAY_KEY, + label='Predefined percentage display format', + component=gr.Dropdown, + component_args=lambda: { + 'choices': tuple(_PREDEFINED_PERCENTAGES_DISPLAY_MAP.keys()), + }, + section=section, + ), + ) + + +class AspectRatioStepScript(scripts.Script): + + def title(self): + return _EXTENSION_NAME + + def show(self, is_img2img): + return scripts.AlwaysVisible + + def ui(self, is_img2img): + if not any([ + opts.arh_show_max_width_or_height, + opts.arh_show_predefined_percentages, + ]): + return # return early as no 'show' options enabled + + with ( + gr.Group(), + gr.Accordion(_EXTENSION_NAME, open=opts.arh_expand_by_default), + contextlib.suppress(AttributeError), + ): + if is_img2img: + inputs = outputs = [self.i2i_w, self.i2i_h] + else: + inputs = outputs = [self.t2i_w, self.t2i_h] + + if opts.arh_show_max_width_or_height: + with gr.Row(): + max_dimension = gr.inputs.Slider( + minimum=_MIN_DIMENSION, + maximum=_MAX_DIMENSION, + step=1, + default=opts.arh_max_width_or_height, + label='Maximum width or height (whichever is higher)', + ) + gr.Button(value='Scale to maximum width or height').click( + fn=_scale_dimensions_to_max_dimension, + inputs=[*inputs, max_dimension], + outputs=outputs, + ) + + if opts.arh_show_predefined_percentages: + display_func = _PREDEFINED_PERCENTAGES_DISPLAY_MAP.get( + opts.arh_predefined_percentages_display_key, + ) + with gr.Column(variant='panel'), gr.Row(variant='compact'): + pps = opts.arh_predefined_percentages + percentages = [ + abs(int(x)) for x in pps.split(',') + ] + for percentage in percentages: + gr.Button(value=display_func(percentage)).click( + fn=partial( + _scale_by_percentage, pct=percentage / 100, + ), + inputs=inputs, + outputs=outputs, + ) + + def after_component(self, component, **kwargs): + element_id = kwargs.get('elem_id') + + if element_id == 'txt2img_width': + self.t2i_w = component + elif element_id == 'txt2img_height': + self.t2i_h = component + elif element_id == 'img2img_width': + self.i2i_w = component + elif element_id == 'img2img_height': + self.i2i_h = component + + +script_callbacks.on_ui_settings(on_ui_settings) diff --git a/scripts/sd_webui_aspect_ratio_helper.py b/scripts/sd_webui_aspect_ratio_helper.py index 5601cba..9fcebbc 100644 --- a/scripts/sd_webui_aspect_ratio_helper.py +++ b/scripts/sd_webui_aspect_ratio_helper.py @@ -1,128 +1,3 @@ -import contextlib -from functools import partial +from aspect_ratio_helper.main import AspectRatioStepScript -import gradio as gr -from modules import script_callbacks -from modules import scripts -from modules import shared -from modules.shared import opts - -from util import _scale_dimensions_to_max_dimension -from util import _scale_by_percentage - -_EXTENSION_NAME = 'Aspect Ratio Helper' - - -def on_ui_settings(): - section = 'aspect_ratio_helper', _EXTENSION_NAME - shared.opts.add_option( - key='arh_expand_by_default', - info=shared.OptionInfo( - default=False, - label='Expand by default', - section=section, - ), - ) - shared.opts.add_option( - key='arh_show_max_width_or_height', - info=shared.OptionInfo( - default=True, - label='Show maximum width or height button', - section=section, - ), - ) - shared.opts.add_option( - key='arh_max_width_or_height', - info=shared.OptionInfo( - default=1024, - label='Maximum width or height default', - section=section, - ), - ) - shared.opts.add_option( - key='arh_show_predefined_percentages', - info=shared.OptionInfo( - default=True, - label='Show percentage buttons', - section=section, - ), - ) - shared.opts.add_option( - key='arh_predefined_percentages', - info=shared.OptionInfo( - default='25, 50, 75, 125, 150, 175, 200', - label='Percentage buttons (75, 125, 150)', - section=section, - ), - ) - - -class AspectRatioStepScript(scripts.Script): - - def title(self): - return _EXTENSION_NAME - - def show(self, is_img2img): - return scripts.AlwaysVisible - - def ui(self, is_img2img): - if not any([ - opts.arh_show_max_width_or_height, - opts.arh_show_predefined_percentages, - ]): - return # return early as no 'show' options enabled - - with ( - gr.Group(), - gr.Accordion(_EXTENSION_NAME, open=opts.arh_expand_by_default), - contextlib.suppress(AttributeError), - ): - if is_img2img: - inputs = outputs = [self.i2i_w, self.i2i_h] - else: - inputs = outputs = [self.t2i_w, self.t2i_h] - - if opts.arh_show_max_width_or_height: - with gr.Row(): - max_dimension = gr.inputs.Slider( - minimum=64, - maximum=2048, - step=16, - default=opts.arh_max_width_or_height, - label='Maximum width or height (whichever is higher)', - ) - gr.Button(value='Scale to maximum width or height').click( - fn=_scale_dimensions_to_max_dimension, - inputs=[*inputs, max_dimension], - outputs=outputs, - ) - - if opts.arh_show_predefined_percentages: - with gr.Column(variant='panel'), gr.Row(variant='compact'): - pps = opts.arh_predefined_percentages - percentages = [ - int(x) for x in pps.split(',') - ] - for percentage in percentages: - gr.Button(value=f'{str(percentage)}%').click( - fn=partial( - _scale_by_percentage, pct=percentage / 100, - ), - inputs=inputs, - outputs=outputs, - ) - - def after_component(self, component, **kwargs): - element_id = kwargs.get('elem_id') - - if element_id == 'txt2img_width': - self.t2i_w = component - elif element_id == 'txt2img_height': - self.t2i_h = component - elif element_id == 'img2img_width': - self.i2i_w = component - elif element_id == 'img2img_height': - self.i2i_h = component - - -script_callbacks.on_ui_settings(on_ui_settings) +__all__ = ['AspectRatioStepScript'] diff --git a/scripts/util.py b/scripts/util.py deleted file mode 100644 index 8fd3153..0000000 --- a/scripts/util.py +++ /dev/null @@ -1,33 +0,0 @@ -_MIN_DIMENSION = 64 -_MAX_DIMENSION = 2048 - - -def _scale_by_percentage(width, height, pct): - aspect_ratio = float(width) / float(height) - step = (pct - 1.0) - new_width = max(int(round(width * (1.0 + step))), 1) - new_height = max(int(round(new_width / aspect_ratio)), 1) - if new_width > _MAX_DIMENSION: - new_width = _MAX_DIMENSION - new_height = max(int(round(new_width / aspect_ratio)), 1) - if new_height > _MAX_DIMENSION: - new_height = _MAX_DIMENSION - new_width = max(int(round(new_height * aspect_ratio)), 1) - if new_width < _MIN_DIMENSION: - new_width = _MIN_DIMENSION - new_height = max(int(round(new_width / aspect_ratio)), 1) - if new_height < _MIN_DIMENSION: - new_height = _MIN_DIMENSION - new_width = max(int(round(new_height * aspect_ratio)), 1) - return new_width, new_height - - -def _scale_dimensions_to_max_dimension(width, height, max_dim): - if not _MIN_DIMENSION < max_dim < _MAX_DIMENSION: - raise ValueError('Invalid dimension provided.') - if max_dim == max(width, height): - return width, height - aspect_ratio = float(width) / float(height) - if width > height: - return max_dim, max(int(round(max_dim / aspect_ratio)), 1) - return max(int(round(max_dim * aspect_ratio)), 1), max_dim diff --git a/tests/_util_test.py b/tests/_util_test.py new file mode 100644 index 0000000..549368b --- /dev/null +++ b/tests/_util_test.py @@ -0,0 +1,131 @@ +import pytest + +from aspect_ratio_helper._util import _display_minus_and_plus +from aspect_ratio_helper._util import _display_multiplication +from aspect_ratio_helper._util import _display_raw_percentage +from aspect_ratio_helper._util import _MAX_DIMENSION +from aspect_ratio_helper._util import _MIN_DIMENSION +from aspect_ratio_helper._util import _scale_by_percentage +from aspect_ratio_helper._util import _scale_dimensions_to_max_dimension + + +@pytest.mark.parametrize( + 'num, expected', + [ + (50, 'x0.5'), + (150, 'x1.5'), + (175, 'x1.75'), + (250, 'x2.5'), + ], +) +def test_display_multiplication(num, expected): + assert _display_multiplication(num) == expected + + +@pytest.mark.parametrize( + 'num, expected', + [ + (50, '50%'), + (75, '75%'), + (100, '100%'), + (150, '150%'), + (250, '250%'), + ], +) +def test_display_raw_percentage(num, expected): + assert _display_raw_percentage(num) == expected + + +@pytest.mark.parametrize( + 'num, expected_output', [ + (150, '+50%'), + (100, '0%'), + (50, '-50%'), + (0, '-100%'), + (200, '+100%'), + (75, '-25%'), + ], +) +def test_display_minus_and_plus(num, expected_output): + assert _display_minus_and_plus(num) == expected_output + + +@pytest.mark.parametrize( + 'width, height, pct, expected', + [ + pytest.param(200, 400, 0.5, (100, 200), id='50_percent_scale_down'), + pytest.param(100, 200, 2.0, (200, 400), id='200_percent_scale_up'), + pytest.param(100, 200, 1.1, (110, 220), id='10_percent_scale_up'), + pytest.param(100, 200, 0.9, (90, 180), id='10_percent_scale_down'), + pytest.param(100, 200, 0.0, (64, 128), id='scale_full_down'), + pytest.param( + _MIN_DIMENSION - 1, + _MIN_DIMENSION - 1, + 0.5, + (_MIN_DIMENSION, _MIN_DIMENSION), + id='scale_below_min_dimension', + ), + pytest.param( + _MAX_DIMENSION + 1, + _MAX_DIMENSION + 1, + 2.0, + (_MAX_DIMENSION, _MAX_DIMENSION), + id='scale_above_max_dimension', + ), + ], +) +def test_scale_by_percentage( + width, height, pct, expected, +): + assert _scale_by_percentage( + width, height, pct, + ) == expected + + +@pytest.mark.parametrize( + 'width, height, max_dim, expected', + [ + pytest.param( + 100, 200, 400, (200, 400), + id='scale_up_to_max_dimension_horizontally', + ), + pytest.param( + 200, 100, 400, (400, 200), + id='scale_up_to_max_dimension_vertically', + ), + pytest.param( + 400, 64, 400, (400, 64), + id='no_scale_up_needed_with_max_dimension_width', + ), + pytest.param( + 64, 400, 400, (64, 400), + id='no_scale_up_needed_with_max_dimension_height', + ), + pytest.param( + _MIN_DIMENSION, _MIN_DIMENSION, _MAX_DIMENSION, + (_MAX_DIMENSION, _MAX_DIMENSION), + id='scale_from_min_to_max', + ), + pytest.param( + _MAX_DIMENSION, _MAX_DIMENSION, _MIN_DIMENSION, + (_MIN_DIMENSION, _MIN_DIMENSION), + id='scale_from_max_to_min', + ), + pytest.param( + 64, 64, _MIN_DIMENSION - 1, + (_MIN_DIMENSION, _MIN_DIMENSION), + id='scale_below_min_dimension_clamps_retains_ar', + ), + pytest.param( + 64, 64, _MAX_DIMENSION + 1, + (_MAX_DIMENSION, _MAX_DIMENSION), + id='scale_above_max_dimension_clamps_retains_ar', + ), + ], +) +def test_scale_dimensions_to_max_dimension( + width, height, max_dim, expected, +): + assert _scale_dimensions_to_max_dimension( + width, height, max_dim, + ) == expected diff --git a/tests/util_test.py b/tests/util_test.py deleted file mode 100644 index f553468..0000000 --- a/tests/util_test.py +++ /dev/null @@ -1,86 +0,0 @@ -import pytest - -from scripts.util import _MIN_DIMENSION -from scripts.util import _MAX_DIMENSION -from scripts.util import _scale_by_percentage -from scripts.util import _scale_dimensions_to_max_dimension - - -@pytest.mark.parametrize( - "width, height, pct, expected", - [ - pytest.param(200, 400, 0.5, (100, 200), id="50_percent_scale_down"), - pytest.param(100, 200, 2.0, (200, 400), id="200_percent_scale_up"), - pytest.param(100, 200, 1.1, (110, 220), id="10_percent_scale_up"), - pytest.param(100, 200, 0.9, (90, 180), id="10_percent_scale_down"), - pytest.param(100, 200, 0.0, (64, 128), id="scale_full_down"), - pytest.param( - _MIN_DIMENSION - 1, - _MIN_DIMENSION - 1, - 0.5, - (_MIN_DIMENSION, _MIN_DIMENSION), - id="scale_below_min_dimension", - ), - pytest.param( - _MAX_DIMENSION + 1, - _MAX_DIMENSION + 1, - 2.0, - (_MAX_DIMENSION, _MAX_DIMENSION), - id="scale_above_max_dimension", - ), - ], -) -def test_scale_by_percentage( - width, height, pct, expected -): - assert _scale_by_percentage( - width, height, pct - ) == expected - - -@pytest.mark.parametrize( - "width, height, max_dim, expected", - [ - pytest.param( - 100, 200, 400, (200, 400), id="scale_up_to_max_dimension_horizontally" - ), - pytest.param( - 200, 100, 400, (400, 200), id="scale_up_to_max_dimension_vertically" - ), - pytest.param( - 400, 64, 400, (400, 64), id="no_scale_up_needed_with_max_dimension_width" - ), - pytest.param( - 64, 400, 400, (64, 400), id="no_scale_up_needed_with_max_dimension_height" - ), - ], -) -def test_scale_dimensions_to_max_dimension( - width, height, max_dim, expected -): - assert _scale_dimensions_to_max_dimension( - width, height, max_dim - ) == expected - - -@pytest.mark.parametrize( - "width, height, max_dim", - [ - pytest.param( - 64, - 64, - _MIN_DIMENSION - 1, - id="scale_below_min_dimension", - ), - pytest.param( - 64, - 64, - _MAX_DIMENSION + 1, - id="scale_above_max_dimension", - ), - - ], -) -def test_error_thrown_given_dim_outside_boundaries(width, height, max_dim): - with pytest.raises(ValueError): - _scale_dimensions_to_max_dimension(width, height, max_dim)