From 03d336acd78298fccb5eac5e126d0fa12d21d5eb Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Fri, 22 Sep 2023 14:39:36 +0200
Subject: [PATCH] Selection grow and feather are now configurable

---
 ai_diffusion/document.py    | 14 +++++++++-----
 ai_diffusion/settings.py    |  8 ++++++++
 ai_diffusion/ui/model.py    |  4 +++-
 ai_diffusion/ui/settings.py | 10 ++++++++--
 4 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/ai_diffusion/document.py b/ai_diffusion/document.py
index e7ef3dc38..2917e4280 100644
--- a/ai_diffusion/document.py
+++ b/ai_diffusion/document.py
@@ -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):
@@ -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)
diff --git a/ai_diffusion/settings.py b/ai_diffusion/settings.py
index c5479d6d8..e44df1da6 100644
--- a/ai_diffusion/settings.py
+++ b/ai_diffusion/settings.py
@@ -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(
diff --git a/ai_diffusion/ui/model.py b/ai_diffusion/ui/model.py
index 86779137d..5e3d21a15 100644
--- a/ai_diffusion/ui/model.py
+++ b/ai_diffusion/ui/model.py
@@ -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)
diff --git a/ai_diffusion/ui/settings.py b/ai_diffusion/ui/settings.py
index a4099893a..6485cc0b6 100644
--- a/ai_diffusion/ui/settings.py
+++ b/ai_diffusion/ui/settings.py
@@ -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()
@@ -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
@@ -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)