Skip to content

Commit

Permalink
Selection grow and feather are now configurable
Browse files Browse the repository at this point in the history
  • Loading branch information
Acly committed Sep 22, 2023
1 parent b0ca44e commit 03d336a
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 8 deletions.
14 changes: 9 additions & 5 deletions ai_diffusion/document.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
from typing import Optional
import krita
from krita import Krita
from .image import Extent, Bounds, Mask, Image
from PyQt5.QtCore import QUuid
from PyQt5.QtGui import QImage

from .image import Extent, Bounds, Mask, Image


class Document:
def __init__(self, krita_document):
Expand Down Expand Up @@ -37,17 +38,20 @@ def check_color_mode(self):
return False, msg_fmt.format("depth", "8-bit integer", depth)
return True, None

def create_mask_from_selection(self):
def create_mask_from_selection(self, grow: float, feather: float):
user_selection = self._doc.selection()
if not user_selection:
return None

selection = user_selection.duplicate()
size_factor = Extent(selection.width(), selection.height()).diagonal
feather_radius = int(0.07 * size_factor)
grow_pixels = int(grow * size_factor)
feather_radius = int(feather * size_factor)

selection.grow(feather_radius, feather_radius)
selection.feather(feather_radius)
if grow_pixels > 0:
selection.grow(grow_pixels, grow_pixels)
if feather_radius > 0:
selection.feather(feather_radius)

bounds = Bounds(selection.x(), selection.y(), selection.width(), selection.height())
bounds = Bounds.pad(bounds, feather_radius, multiple=8)
Expand Down
8 changes: 8 additions & 0 deletions ai_diffusion/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,14 @@ class Settings:
"Server Arguments", "", "Additional command line arguments passed to the server"
)

_selection_grow = Setting(
"Selection Grow", 7, "Selection area is expanded by a fraction of its size"
)

_selection_feather = Setting(
"Selection Feather", 7, "The border is blurred by a fraction of selection size"
)

_fixed_seed = Setting("Use Fixed Seed", False, "Fixes the random seed to a specific value")

_random_seed = Setting(
Expand Down
4 changes: 3 additions & 1 deletion ai_diffusion/ui/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,9 @@ def generate(self):
image = None
extent = self._doc.extent

mask = self._doc.create_mask_from_selection()
mask = self._doc.create_mask_from_selection(
grow=settings.selection_grow / 100, feather=settings.selection_feather / 100
)
image_bounds = workflow.compute_bounds(extent, mask.bounds if mask else None, self.strength)
if mask is not None or self.strength < 1.0:
image = self._doc.get_image(image_bounds, exclude_layer=self._layer)
Expand Down
10 changes: 8 additions & 2 deletions ai_diffusion/ui/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,8 +95,9 @@ def value(self, v):


class SliderSetting(SettingWidget):
def __init__(self, setting: Setting, parent=None, minimum=0, maximum=100):
def __init__(self, setting: Setting, parent=None, minimum=0, maximum=100, format="{}"):
super().__init__(setting, parent)
self._format_string = format

slider_widget = QWidget(self)
slider_layout = QHBoxLayout()
Expand All @@ -115,7 +116,7 @@ def __init__(self, setting: Setting, parent=None, minimum=0, maximum=100):
self._layout.addWidget(slider_widget, alignment=Qt.AlignRight)

def _change_value(self, value: int):
self._label.setText(str(value))
self._label.setText(self._format_string.format(value))
self.value_changed.emit()

@property
Expand Down Expand Up @@ -697,6 +698,11 @@ class DiffusionSettings(SettingsTab):
def __init__(self):
super().__init__("Diffusion Settings")

self.add("selection_grow", SliderSetting(Settings._selection_grow, self, 0, 25, "{} %"))
self.add(
"selection_feather", SliderSetting(Settings._selection_feather, self, 0, 25, "{} %")
)

self.add("random_seed", TextSetting(Settings._random_seed, self))
self._fixed_seed_checkbox = QCheckBox("Use fixed seed", self)
self._fixed_seed_checkbox.stateChanged.connect(self.write)
Expand Down

0 comments on commit 03d336a

Please sign in to comment.