Try more robust method of inserting the recording and restoring activ…

…e layer

ai_diffusion/document.py

@@ -42,7 +42,12 @@ def get_layer_image(self, layer: krita.Node, bounds: Bounds | None) -> Image:
         raise NotImplementedError
 
     def insert_layer(
-        self, name: str, img: Image | None = None, bounds: Bounds | None = None, make_active=True
+        self,
+        name: str,
+        img: Image | None = None,
+        bounds: Bounds | None = None,
+        make_active=True,
+        below: krita.Node | None = None,
     ) -> krita.Node:
         raise NotImplementedError
 
@@ -233,11 +238,16 @@ def get_layer_image(self, layer: krita.Node, bounds: Bounds | None):
         return Image(QImage(data, *bounds.extent, QImage.Format.Format_ARGB32))
 
     def insert_layer(
-        self, name: str, img: Image | None = None, bounds: Bounds | None = None, make_active=True
+        self,
+        name: str,
+        img: Image | None = None,
+        bounds: Bounds | None = None,
+        make_active=True,
+        below: krita.Node | None = None,
     ):
         with RestoreActiveLayer(self) if not make_active else nullcontext():
             layer = self._doc.createNode(name, "paintlayer")
-            self._doc.rootNode().addChildNode(layer, None)
+            self._doc.rootNode().addChildNode(layer, _find_layer_above(self._doc, below))
             if img and bounds:
                 layer.setPixelData(img.data, *bounds)
                 self._doc.refreshProjection()
@@ -346,6 +356,15 @@ 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())
 
@@ -377,6 +396,11 @@ def __exit__(self, exc_type, exc_value, traceback):
 
     async def _restore(self):
         if self.layer:
+            if self.layer == self.document.active_layer:
+                # Maybe whatever event we expected to change the active layer hasn't happened yet.
+                await eventloop.wait_until(
+                    lambda: self.document.active_layer != self.layer, no_error=True
+                )
             self.document.active_layer = self.layer
 
 

ai_diffusion/eventloop.py

@@ -47,10 +47,10 @@ def stop():
         pass
 
 
-async def wait_until(condition: Callable[[], bool], timeout=1.0):
-    start = time.time()
-    while not condition():
-        if time.time() - start > timeout:
+async def wait_until(condition: Callable[[], bool], iterations=10, no_error=False):
+    while not condition() and iterations > 0:
+        iterations -= 1
+        if iterations == 0 and not no_error:
             raise TimeoutError("Timeout while waiting for action to complete")
         await asyncio.sleep(0.01)
 

ai_diffusion/model.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from contextlib import nullcontext
 from pathlib import Path
 from enum import Enum
 from typing import NamedTuple
@@ -469,7 +470,9 @@ def toggle_record(self, active: bool):
             self._is_recording = active
             self.is_active = active
             self.is_recording_changed.emit(active)
-            if not active:
+            if active:
+                self._add_recording_layer()
+            else:
                 self._insert_frames()
 
     def handle_job_finished(self, job: Job):
@@ -500,21 +503,28 @@ def set_result(self, value: Image, bounds: Bounds):
         if self.is_recording:
             self._keyframes.append((value, bounds))
 
-    def _insert_frames(self):
-        if len(self._keyframes) == 0:
-            return
+    def _add_recording_layer(self, restore_active=True):
         doc = self._model.document
         if self._recording_layer and self._recording_layer.parentNode() is None:
             self._recording_layer = None
         if self._recording_layer is None:
-            self._recording_layer = doc.insert_layer(
-                f"[Recording] {self._model.prompt}", make_active=False
-            )
-            self._recording_layer.enableAnimation()
-            self._recording_layer.setPinnedToTimeline(True)
+            with RestoreActiveLayer(doc) if restore_active else nullcontext():
+                self._recording_layer = doc.insert_layer(
+                    f"[Recording] {self._model.prompt}", below=doc.active_layer
+                )
+                self._recording_layer.enableAnimation()
+                self._recording_layer.setPinnedToTimeline(True)
+        return self._recording_layer
 
-        eventloop.run(self._insert_frames_into_timeline(self._recording_layer, self._keyframes))
-        self._keyframes = []
+    def _insert_frames(self):
+        if len(self._keyframes) > 0:
+            layer = self._add_recording_layer(restore_active=False)
+            eventloop.run(
+                _report_errors(
+                    self._model, self._insert_frames_into_timeline(layer, self._keyframes)
+                )
+            )
+            self._keyframes = []
 
     async def _insert_frames_into_timeline(
         self, layer: krita.Node, frames: list[tuple[Image, Bounds]]
@@ -537,7 +547,7 @@ async def _insert_frames_into_timeline(
         doc.end_time = doc.current_time
 
 
-async def _report_errors(parent, coro):
+async def _report_errors(parent: Model, coro):
     try:
         return await coro
     except NetworkError as e: