From a76ddf16ab4882a1e27ad1f579162285e49bdd77 Mon Sep 17 00:00:00 2001
From: Acly <aclysia@gmail.com>
Date: Sat, 23 Sep 2023 11:24:23 +0200
Subject: [PATCH] Support group layers as control image #10 * list nested
 layers * list layers in the order of the layer stack

---
 ai_diffusion/document.py  | 19 ++++++++++++++-----
 ai_diffusion/ui/widget.py |  6 ++++--
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/ai_diffusion/document.py b/ai_diffusion/document.py
index 2917e4280..03ab5661d 100644
--- a/ai_diffusion/document.py
+++ b/ai_diffusion/document.py
@@ -1,7 +1,7 @@
 from typing import Optional
 import krita
 from krita import Krita
-from PyQt5.QtCore import QUuid
+from PyQt5.QtCore import QUuid, QByteArray
 from PyQt5.QtGui import QImage
 
 from .image import Extent, Bounds, Mask, Image
@@ -76,7 +76,9 @@ def get_image(self, bounds: Bounds = None, exclude_layer=None):
         return Image(img)
 
     def get_layer_image(self, layer, bounds: Bounds):
-        return Image(QImage(layer.pixelData(*bounds), *bounds.extent, QImage.Format_ARGB32))
+        data: QByteArray = layer.projectionPixelData(*bounds)
+        assert data is not None and data.size() >= bounds.extent.pixel_count * 4
+        return Image(QImage(data, *bounds.extent, QImage.Format_ARGB32))
 
     def insert_layer(
         self, name: str, img: Image, bounds: Bounds, below: Optional[krita.Node] = None
@@ -110,12 +112,19 @@ def hide_layer(self, layer):
         return layer
 
     @property
-    def paint_layers(self):
-        return [node for node in self._doc.rootNode().childNodes() if node.type() == "paintlayer"]
+    def image_layers(self):
+        return list(_traverse_layers(self._doc.rootNode(), ["paintlayer", "grouplayer"]))
 
     def find_layer(self, id: QUuid):
-        return next((layer for layer in self.paint_layers if layer.uniqueId() == id), None)
+        return next((layer for layer in self.image_layers if layer.uniqueId() == id), None)
 
     @property
     def active_layer(self):
         return self._doc.activeNode()
+
+
+def _traverse_layers(node, type_filter=None):
+    for child in node.childNodes():
+        yield from _traverse_layers(child, type_filter)
+        if not type_filter or child.type() in type_filter:
+            yield child
diff --git a/ai_diffusion/ui/widget.py b/ai_diffusion/ui/widget.py
index 82acc8746..7fbed9b9e 100644
--- a/ai_diffusion/ui/widget.py
+++ b/ai_diffusion/ui/widget.py
@@ -151,11 +151,13 @@ def _notify(self):
             self.changed.emit()
 
     def update_and_select_layer(self, id: QUuid):
-        layers = self._model.document.paint_layers
+        layers = reversed(self._model.document.image_layers)
         self.layer_select.clear()
+        index = -1
         for layer in layers:
             self.layer_select.addItem(layer.name(), layer.uniqueId())
-        index = next((i for i, layer in enumerate(layers) if layer.uniqueId() == id), -1)
+            if layer.uniqueId() == id:
+                index = self.layer_select.count() - 1
         if index == -1:
             self.remove_button.click()
         else: