Adds image-to-image support for Z-Image

Author: w4ffl35

src/airunner/components/art/gui/widgets/canvas/image_manipulation_tools_container.py

@@ -1,9 +1,17 @@
-from PySide6.QtCore import Slot
+from PySide6.QtCore import Slot, QTimer
 from airunner.components.art.gui.widgets.canvas.templates.image_manipulation_tools_container_ui import (
     Ui_image_manipulation_tools_container,
 )
 from airunner.components.application.gui.widgets.base_widget import BaseWidget
 from airunner.utils.settings.get_qsettings import get_qsettings
+from airunner.enums import StableDiffusionVersion
+
+
+# Versions that don't support ControlNet or Inpaint
+_NO_CONTROLNET_INPAINT_VERSIONS = (
+    StableDiffusionVersion.Z_IMAGE_TURBO.value,
+    StableDiffusionVersion.Z_IMAGE_BASE.value,
+)
 
 
 class ImageManipulationToolsContainer(BaseWidget):
@@ -12,23 +20,82 @@ class ImageManipulationToolsContainer(BaseWidget):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.qsettings = get_qsettings()
+        self._controlnet_tab = None
+        self._inpaint_tab = None
+        self._last_version = None
         self.ui.image_manipulation_tools_tab_container.currentChanged.connect(
             self.on_tab_changed
         )
+        # Set up a timer to periodically check version changes
+        self._version_check_timer = QTimer(self)
+        self._version_check_timer.setInterval(500)  # Check every 500ms
+        self._version_check_timer.timeout.connect(self._check_version_and_update_tabs)
 
     @Slot(int)
     def on_tab_changed(self, index: int):
         self.qsettings.setValue(
             "tabs/image_manipulation_tools_container/active_index", index
         )
 
+    def _check_version_and_update_tabs(self):
+        """Check if version changed and update tab visibility accordingly."""
+        current_version = self.generator_settings.version
+        if current_version != self._last_version:
+            self._last_version = current_version
+            self._update_tab_visibility()
+
+    def _update_tab_visibility(self):
+        """Hide/show ControlNet and Inpaint tabs based on model version.
+        
+        Z-Image models don't support ControlNet or Inpaint pipelines,
+        so those tabs should be hidden when Z-Image is selected.
+        """
+        tab_widget = self.ui.image_manipulation_tools_tab_container
+        current_version = self.generator_settings.version
+        should_hide = current_version in _NO_CONTROLNET_INPAINT_VERSIONS
+        
+        # Store references to removed tabs so they can be restored
+        if should_hide:
+            # Remove ControlNet tab (index 1) and Inpaint tab (index 2)
+            # Note: Remove in reverse order to preserve indices
+            if tab_widget.count() > 2:
+                self._inpaint_tab = tab_widget.widget(2)
+                tab_widget.removeTab(2)
+            if tab_widget.count() > 1:
+                self._controlnet_tab = tab_widget.widget(1)
+                tab_widget.removeTab(1)
+        else:
+            # Restore tabs if they were previously removed
+            if self._controlnet_tab is not None and tab_widget.count() == 1:
+                tab_widget.insertTab(1, self._controlnet_tab, "Controlnet")
+                self._controlnet_tab = None
+            if self._inpaint_tab is not None and tab_widget.count() == 2:
+                tab_widget.insertTab(2, self._inpaint_tab, "Inpaint")
+                self._inpaint_tab = None
+
     def showEvent(self, event):
         super().showEvent(event)
+        # Update tab visibility based on current model version
+        self._last_version = self.generator_settings.version
+        self._update_tab_visibility()
+        
+        # Start the version check timer
+        self._version_check_timer.start()
+        
+        # Restore active tab index
         active_index = int(
             self.qsettings.value(
                 "tabs/image_manipulation_tools_container/active_index", 0
             )
         )
+        # Ensure index is valid for current tab count
+        max_index = self.ui.image_manipulation_tools_tab_container.count() - 1
+        active_index = min(active_index, max_index)
         self.ui.image_manipulation_tools_tab_container.setCurrentIndex(
             active_index
         )
+
+    def hideEvent(self, event):
+        super().hideEvent(event)
+        # Stop the timer when widget is hidden
+        self._version_check_timer.stop()

src/airunner/components/art/gui/widgets/canvas/input_image.py

@@ -377,7 +377,9 @@ def load_image_from_object(self, image: Image):
 
             # Adjust scene rect to the item bounds if present
             if self._scene.item and hasattr(self._scene.item, "boundingRect"):
-                rect = self._scene.item.boundingRect()
+                # Use sceneBoundingRect() to get the rect in scene coordinates
+                # (includes the item's position)
+                rect = self._scene.item.sceneBoundingRect()
                 self._scene.setSceneRect(rect)
 
             # Force scene and view updates

src/airunner/components/art/gui/widgets/canvas/input_image_scene.py

@@ -1,6 +1,7 @@
 from typing import Dict
 from PIL.Image import Image
 from PIL import ImageQt
+from PySide6.QtCore import QPoint
 
 from airunner.components.art.data.controlnet_settings import ControlnetSettings
 from airunner.components.art.data.drawingpad_settings import DrawingPadSettings
@@ -29,6 +30,37 @@ def __init__(
     def settings_key(self):
         return self._settings_key
 
+    def _get_default_image_position(self) -> QPoint:
+        """Override to always place images at (0, 0) in input image scenes.
+        
+        Unlike the main canvas which uses drawing_pad_settings for position,
+        input image scenes are small preview panels that should always display
+        images starting at the origin.
+        
+        Returns:
+            QPoint at (0, 0).
+        """
+        return QPoint(0, 0)
+
+    def _update_item_position(
+        self, root_point: QPoint, canvas_offset
+    ) -> None:
+        """Override to always position items at (0, 0) in input image scenes.
+        
+        Input image scenes are small preview panels that should always display
+        images at the origin, ignoring any stored position or canvas offset.
+        
+        Args:
+            root_point: Ignored - always uses (0, 0).
+            canvas_offset: Ignored - always uses (0, 0).
+        """
+        try:
+            if self.item is not None:
+                self.item.setPos(0, 0)
+        except (RuntimeError, AttributeError):
+            # Item was deleted or is no longer valid
+            pass
+
     @property
     def current_active_image(self) -> Image:
         if self._is_mask:

src/airunner/components/art/managers/stablediffusion/mixins/sd_generation_preparation_mixin.py

@@ -113,7 +113,9 @@ def _prepare_data(self, active_rect=None) -> Dict:
             data.update({"width": width, "height": height})
 
         if self.is_img2img:
-            image = self.img2img_image
+            # Use image from image_request if available (passed from GUI),
+            # otherwise fall back to img2img_image property
+            image = self.image_request.image if self.image_request.image is not None else self.img2img_image
             if (
                 data["num_inference_steps"]
                 < AIRUNNER_MIN_NUM_INFERENCE_STEPS_IMG2IMG