improves image input controls

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="airunner",
-    version="3.0.12",
+    version="3.0.13",
     author="Capsize LLC",
     description="A Stable Diffusion GUI",
     long_description=open("README.md", "r", encoding="utf-8").read(),

src/airunner/alembic/versions/6d58cbfd61fb_add_lock_input_image_column_to_image_to_.py

@@ -0,0 +1,34 @@
+"""Add lock_input_image column to image_to_image_settings
+
+Revision ID: 6d58cbfd61fb
+Revises: 72d9134823cb
+Create Date: 2024-10-09 10:05:11.349862
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import sqlite
+
+from airunner.utils.db.column_exists import column_exists
+
+# revision identifiers, used by Alembic.
+revision: str = '6d58cbfd61fb'
+down_revision: Union[str, None] = '72d9134823cb'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade():
+    if not column_exists('image_to_image_settings', 'lock_input_image'):
+        op.add_column('image_to_image_settings', sa.Column('lock_input_image', sa.Boolean, default=False))
+    if not column_exists('controlnet_settings', 'lock_input_image'):
+        op.add_column('controlnet_settings', sa.Column('lock_input_image', sa.Boolean, default=False))
+
+def downgrade():
+    if column_exists('image_to_image_settings', 'lock_input_image'):
+        op.drop_column('image_to_image_settings', 'lock_input_image')
+    if column_exists('controlnet_settings', 'lock_input_image'):
+        op.drop_column('image_to_image_settings', 'lock_input_image')
+    # ### end Alembic commands ###

src/airunner/data/models/settings_models.py

@@ -91,6 +91,7 @@ class ControlnetSettings(Base):
     conditioning_scale = Column(Integer, default=100)
     guidance_scale = Column(Integer, default=750)
     controlnet = Column(String, default="Canny")
+    lock_input_image = Column(Boolean, default=False)
 
 
 class ImageToImageSettings(Base):
@@ -99,6 +100,7 @@ class ImageToImageSettings(Base):
     image = Column(String, nullable=True)
     enabled = Column(Boolean, default=False)
     use_grid_image_as_input = Column(Boolean, default=False)
+    lock_input_image = Column(Boolean, default=False)
 
 
 class OutpaintSettings(Base):

src/airunner/enums.py

@@ -242,6 +242,7 @@ class SignalCode(Enum):
     MASK_LAYER_TOGGLED = enum.auto()
     MASK_UPDATED = enum.auto()
     HISTORY_UPDATED = enum.auto()
+    CANVAS_IMAGE_UPDATED_SIGNAL = enum.auto()
 
 class EngineResponseCode(Enum):
     STATUS = 100

src/airunner/styles/dark_theme/styles.qss

@@ -730,3 +730,7 @@ QMenu, QToolBar, QMenuBar {
 QMenu {
     padding: 10px;
 }
+
+SwitchWidget:enabled[checked="true"] {
+    background-color: rgba(0, 132, 185, 100);
+}

src/airunner/utils/db/column_exists.py

@@ -0,0 +1,8 @@
+from alembic import op
+import sqlalchemy as sa
+
+def column_exists(table_name, column_name):
+    conn = op.get_bind()
+    inspector = sa.inspect(conn)
+    columns = [col['name'] for col in inspector.get_columns(table_name)]
+    return column_name in columns

src/airunner/widgets/canvas/brush_scene.py

@@ -217,6 +217,7 @@ def _handle_left_mouse_release(self, event) -> bool:
                 self.emit_signal(SignalCode.GENERATE_MASK)
         session.commit()
         session.close()
+        self.emit_signal(SignalCode.CANVAS_IMAGE_UPDATED_SIGNAL)
         if self.drawing_pad_settings.mask_layer_enabled:
             self.initialize_image()
             self.emit_signal(SignalCode.MASK_UPDATED)

src/airunner/widgets/canvas/custom_scene.py

@@ -117,8 +117,11 @@ def current_active_image(self) -> Image:
 
     @current_active_image.setter
     def current_active_image(self, image: Image):
-        self._update_current_settings("image", convert_image_to_base64(image))
-        self.initialize_image(image)
+        if image is not None:
+            image = convert_image_to_base64(image)
+        self._update_current_settings("image", image)
+        if self.settings_key == "drawing_pad_settings":
+            self.emit_signal(SignalCode.CANVAS_IMAGE_UPDATED_SIGNAL)
 
     @image_pivot_point.setter
     def image_pivot_point(self, value):
@@ -173,7 +176,7 @@ def on_paste_image_from_clipboard(self):
         if self.application_settings.resize_on_paste:
             image = self._resize_image(image)
         image = convert_image_to_base64(image)
-        self._update_current_settings("image", image)
+        self.current_active_image = image
         self.refresh_image(self.current_active_image)
 
     def on_load_image_from_path(self, message):
@@ -189,6 +192,7 @@ def on_load_image_signal(self, image_path: str):
         if self.application_settings.resize_on_paste:
             image = self._resize_image(image)
         self.current_active_image = image
+        self.initialize_image(image)
 
     def on_apply_filter_signal(self, message):
         self._apply_filter(message)
@@ -230,7 +234,7 @@ def on_image_generated_signal(self, response):
             callback(response)
 
     def on_canvas_clear_signal(self):
-        self._update_current_settings("image", None)
+        self.current_active_image = None
         self.delete_image()
         self._clear_history()
 
@@ -269,7 +273,7 @@ def _history_set_image(self, data: dict):
                 self.delete_image()
             else:
                 self.current_active_image = data["image"]
-                self.initialize_image()
+                self.initialize_image(self.current_active_image)
 
     def showEvent(self, event):
         super().showEvent(event)
@@ -411,7 +415,7 @@ def delete_image(self):
 
         if self.painter and self.painter.isActive():
             self.painter.end()
-        self._update_current_settings("image", None)
+        self.current_active_image = None
         self.image = None
         self.initialize_image()
 
@@ -632,8 +636,7 @@ def _add_image_to_scene(
                 action=GeneratorSection.OUTPAINT.value
             )
         # self._set_current_active_image(image)
-        base64_image = convert_image_to_base64(image)
-        self._update_current_settings("image", base64_image)
+        self.current_active_image = image
         q_image = ImageQt.ImageQt(image)
         self.item.setPixmap(QPixmap.fromImage(q_image))
         self.item.setZValue(0)
@@ -712,7 +715,7 @@ def rotate_image(
         if image is not None:
             self._add_image_to_undo()
             image = image.rotate(angle, expand=True)
-            self._update_current_settings("image", convert_image_to_base64(image))
+            self.current_active_image = image
             self.initialize_image(image)
 
     def _add_undo_history(self, data: dict):

src/airunner/widgets/canvas/input_image.py

@@ -7,6 +7,7 @@
 from PySide6.QtGui import QPixmap, QImage, Qt, QPen
 from PySide6.QtWidgets import QGraphicsScene
 
+from airunner.enums import SignalCode
 from airunner.settings import VALID_IMAGE_FILES
 from airunner.utils.convert_base64_to_image import convert_base64_to_image
 from airunner.utils.convert_image_to_base64 import convert_image_to_base64
@@ -69,32 +70,50 @@ def showEvent(self, event):
         else:
             self.ui.mask_blur_slider_widget.hide()
 
-        self.ui.enable_checkbox.blockSignals(True)
-        self.ui.use_grid_image_as_input_checkbox.blockSignals(True)
-        self.ui.enable_checkbox.setChecked(self.current_settings.enabled)
+        self.ui.EnableSwitch.toggled.connect(self.enabled_toggled)
         if self.settings_key == "outpaint_settings":
             self.ui.import_button.hide()
-            self.ui.use_grid_image_as_input_checkbox.hide()
-            self.ui.use_grid_image_as_input_checkbox.hide()
+            self.ui.link_to_grid_image_button.hide()
+            self.ui.link_to_grid_image_button.hide()
+            self.ui.lock_input_image_button.hide()
         else:
-            self.ui.use_grid_image_as_input_checkbox.setChecked(
+            self.ui.EnableSwitch.blockSignals(True)
+            self.ui.link_to_grid_image_button.blockSignals(True)
+            self.ui.lock_input_image_button.blockSignals(True)
+            self.ui.EnableSwitch.checked = self.current_settings.enabled
+            self.ui.EnableSwitch.setChecked(self.current_settings.enabled)
+            self.ui.link_to_grid_image_button.setChecked(
                 self.current_settings.use_grid_image_as_input
             )
-        self.ui.enable_checkbox.blockSignals(False)
-        self.ui.use_grid_image_as_input_checkbox.blockSignals(False)
+            self.ui.lock_input_image_button.setChecked(
+                self.current_settings.lock_input_image or False  # Provide a default value
+            )
+            self.ui.EnableSwitch.blockSignals(False)
+            self.ui.link_to_grid_image_button.blockSignals(False)
+            self.ui.lock_input_image_button.blockSignals(False)
+
+            if self.current_settings.use_grid_image_as_input:
+                self.load_image_from_grid()
+                return
         self.load_image_from_settings()
 
     @Slot(bool)
     def enabled_toggled(self, val):
         self.update_current_settings("enabled", val)
 
+    @Slot(bool)
+    def lock_input_image(self, val):
+        self.update_current_settings("lock_input_image", val)
+
+    @Slot(bool)
+    def refresh_input_image_from_grid(self):
+        self.load_image_from_grid(forced=True)
+
     @Slot(bool)
     def use_grid_image_as_input_toggled(self, val):
         self.update_current_settings("use_grid_image_as_input", val)
         if val is True:
-            base64_image = self.drawing_pad_settings.image
-            self.update_current_settings("image", base64_image)
-            self.load_image_from_settings()
+            self.load_image_from_grid()
 
     @Slot()
     def import_clicked(self):
@@ -123,6 +142,16 @@ def load_image(self, file_path: str):
         if image is not None:
             self.update_current_settings("image", convert_image_to_base64(image))
 
+    def load_image_from_grid(self, forced=False):
+        if self.settings_key == "outpaint_settings":
+            return
+        if not forced and not self.current_settings.use_grid_image_as_input:
+            return
+        if not forced and self.current_settings.lock_input_image:
+            return
+        self.update_current_settings("image", self.drawing_pad_settings.image)
+        self.load_image_from_settings()
+
     def load_image_from_settings(self):
         if self.use_generated_image:
             image = self.current_settings.generated_image
@@ -136,8 +165,8 @@ def load_image_from_settings(self):
             image = convert_base64_to_image(image)
             if image is not None:
                 self.load_image_from_object(image)
-            else:
-                print("image is none")
+                return
+        self.ui.image_container.setScene(None)
 
     def load_image_from_object(self, image: Image):
         if image is None:

src/airunner/widgets/canvas/input_image_container.py

@@ -10,12 +10,18 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.register(SignalCode.MASK_GENERATOR_WORKER_RESPONSE_SIGNAL, self.on_mask_generator_worker_response_signal)
         self.register(SignalCode.MASK_UPDATED, self.on_mask_generator_worker_response_signal)
+        self.register(SignalCode.CANVAS_IMAGE_UPDATED_SIGNAL, self.on_load_image_from_grid_signal)
         self.input_image = None
+        self.generated_image = None
 
     def on_mask_generator_worker_response_signal(self, message):
         if self.input_image:
             self.input_image.on_mask_generator_worker_response_signal()
 
+    def on_load_image_from_grid_signal(self):
+        if self.input_image:
+            self.input_image.load_image_from_grid()
+
     def showEvent(self, event):
         if self.input_image is None:
             settings_key = self.settings_key
@@ -24,8 +30,8 @@ def showEvent(self, event):
             label = "Image-to-Image"
             if settings_key == "controlnet_settings":
                 label = "Controlnet"
-                self.input_image = InputImage(settings_key=self.settings_key, use_generated_image=True)
-                self.ui.tabWidget.addTab(self.input_image, "Generated Image")
+                self.generated_image = InputImage(settings_key=self.settings_key, use_generated_image=True)
+                self.ui.tabWidget.addTab(self.generated_image, "Generated Image")
             elif settings_key == "outpaint_settings":
                 label = "Inpaint / Outpaint"
             self.ui.label.setText(label)