Changes to layer insertion/update behavior:

- Always insert preview layer at top
- When selecting a thumbnail to preview, move the layer to the top
- Don't change the active layer when inserting/updating the preview
- Do change the active layer when applying a result
- For generate control layers or live/upscaling: always insert on top and make active (like apply)

ai_diffusion/document.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from contextlib import nullcontext
 from typing import cast
 import krita
 from krita import Krita
@@ -7,7 +8,7 @@
 
 from .image import Extent, Bounds, Mask, Image
 from .pose import Pose
-from .util import client_logger as log
+from . import eventloop
 
 
 class Document:
@@ -44,14 +45,10 @@ def get_image(
     def get_layer_image(self, layer: krita.Node, bounds: Bounds | None) -> Image:
         raise NotImplementedError
 
-    def insert_layer(
-        self, name: str, img: Image, bounds: Bounds, below: krita.Node | None = None
-    ) -> krita.Node:
+    def insert_layer(self, name: str, img: Image, bounds: Bounds, make_active=True) -> krita.Node:
         raise NotImplementedError
 
-    def insert_vector_layer(
-        self, name: str, svg: str, below: krita.Node | None = None
-    ) -> krita.Node:
+    def insert_vector_layer(self, name: str, svg: str) -> krita.Node:
         raise NotImplementedError
 
     def set_layer_content(self, layer: krita.Node, img: Image, bounds: Bounds):
@@ -60,6 +57,9 @@ def set_layer_content(self, layer: krita.Node, img: Image, bounds: Bounds):
     def hide_layer(self, layer: krita.Node):
         raise NotImplementedError
 
+    def move_to_top(self, layer: krita.Node):
+        raise NotImplementedError
+
     def resize(self, extent: Extent):
         raise NotImplementedError
 
@@ -73,6 +73,10 @@ def create_layer_observer(self) -> LayerObserver:
     def active_layer(self) -> krita.Node:
         raise NotImplementedError
 
+    @active_layer.setter
+    def active_layer(self, layer: krita.Node):
+        pass
+
     @property
     def resolution(self):
         return 0
@@ -192,18 +196,17 @@ def get_layer_image(self, layer: krita.Node, bounds: Bounds | None):
         assert data is not None and data.size() >= bounds.extent.pixel_count * 4
         return Image(QImage(data, *bounds.extent, QImage.Format.Format_ARGB32))
 
-    def insert_layer(self, name: str, img: Image, bounds: Bounds, below: krita.Node | None = None):
-        layer = self._doc.createNode(name, "paintlayer")
-        above = _find_layer_above(self._doc, below)
-        self._doc.rootNode().addChildNode(layer, above)
-        layer.setPixelData(img.data, *bounds)
-        self._doc.refreshProjection()
-        return layer
+    def insert_layer(self, name: str, img: Image, bounds: Bounds, make_active=True):
+        with RestoreActiveLayer(self._doc) if not make_active else nullcontext():
+            layer = self._doc.createNode(name, "paintlayer")
+            self._doc.rootNode().addChildNode(layer, None)
+            layer.setPixelData(img.data, *bounds)
+            self._doc.refreshProjection()
+            return layer
 
-    def insert_vector_layer(self, name: str, svg: str, below: krita.Node | None = None):
+    def insert_vector_layer(self, name: str, svg: str):
         layer = self._doc.createVectorLayer(name)
-        above = _find_layer_above(self._doc, below)
-        self._doc.rootNode().addChildNode(layer, above)
+        self._doc.rootNode().addChildNode(layer, None)
         layer.addShapesFromSvg(svg)
         self._doc.refreshProjection()
         return layer
@@ -224,6 +227,14 @@ def hide_layer(self, layer: krita.Node):
         self._doc.refreshProjection()
         return layer
 
+    def move_to_top(self, layer: krita.Node):
+        parent = layer.parentNode()
+        if parent.childNodes()[-1] == layer:
+            return  # already top-most layer
+        with RestoreActiveLayer(self._doc):
+            parent.removeChildNode(layer)
+            parent.addChildNode(layer, None)
+
     def resize(self, extent: Extent):
         res = self._doc.resolution()
         self._doc.scaleImage(extent.width, extent.height, res, res, "Bilinear")
@@ -239,6 +250,10 @@ def create_layer_observer(self):
     def active_layer(self):
         return self._doc.activeNode()
 
+    @active_layer.setter
+    def active_layer(self, layer: krita.Node):
+        self._doc.setActiveNode(layer)
+
     @property
     def resolution(self):
         return self._doc.resolution() / 72.0  # KisImage::xRes which is applied to vectors
@@ -251,15 +266,6 @@ def _traverse_layers(node: krita.Node, type_filter=None):
             yield child
 
 
-def _find_layer_above(doc: krita.Document, layer_below: krita.Node | None):
-    if layer_below:
-        nodes = doc.rootNode().childNodes()
-        index = nodes.index(layer_below)
-        if index >= 1:
-            return nodes[index - 1]
-    return None
-
-
 def _selection_bounds(selection: krita.Selection):
     return Bounds(selection.x(), selection.y(), selection.width(), selection.height())
 
@@ -275,6 +281,25 @@ def _selection_is_entire_document(selection: krita.Selection, extent: Extent):
     return is_opaque
 
 
+class RestoreActiveLayer:
+    layer: krita.Node | None = None
+
+    def __init__(self, document: krita.Document):
+        self.document = document
+
+    def __enter__(self):
+        self.layer = self.document.activeNode()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Some operations like inserting a new layer change the active layer as a side effect.
+        # It doesn't happen directly, so changing it back in the same call doesn't work.
+        eventloop.run(self._restore())
+
+    async def _restore(self):
+        if self.layer:
+            self.document.setActiveNode(self.layer)
+
+
 class LayerObserver(QObject):
     managed_layer_types = [
         "paintlayer",

ai_diffusion/model.py

@@ -287,8 +287,11 @@ def show_preview(self, job_id: str, index: int, name_prefix="Preview"):
         if self._layer is not None:
             self._layer.setName(name)
             self._doc.set_layer_content(self._layer, job.results[index], job.bounds)
+            self._doc.move_to_top(self._layer)
         else:
-            self._layer = self._doc.insert_layer(name, job.results[index], job.bounds)
+            self._layer = self._doc.insert_layer(
+                name, job.results[index], job.bounds, make_active=False
+            )
             self._layer.setLocked(True)
 
     def hide_preview(self):
@@ -299,6 +302,7 @@ def apply_result(self):
         assert self._layer and self.can_apply_result
         self._layer.setLocked(False)
         self._layer.setName(self._layer.name().replace("[Preview]", "[Generated]"))
+        self._doc.active_layer = self._layer
         self._layer = None
         self.jobs.selection = None
 
@@ -307,9 +311,9 @@ def add_control_layer(self, job: Job, result: dict | None):
         if job.control.mode is ControlMode.pose and result is not None:
             pose = Pose.from_open_pose_json(result)
             pose.scale(job.bounds.extent)
-            return self._doc.insert_vector_layer(job.prompt, pose.to_svg(), below=self._layer)
+            return self._doc.insert_vector_layer(job.prompt, pose.to_svg())
         elif len(job.results) > 0:
-            return self._doc.insert_layer(job.prompt, job.results[0], job.bounds, below=self._layer)
+            return self._doc.insert_layer(job.prompt, job.results[0], job.bounds)
         return self.document.active_layer  # Execution was cached and no image was produced
 
     def add_upscale_layer(self, job: Job):