Fix visibility and enabled state of "Add pose" button #200

Acly · Dec 9, 2023 · 102c2db · 102c2db
1 parent d7554ae
commit 102c2db
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/ai_diffusion/control.py b/ai_diffusion/control.py
@@ -56,7 +56,7 @@ def __init__(self, model: model.Model, mode: ControlMode, layer_id: QUuid):
 
     @property
     def layer(self):
-        layer = self._model.image_layers.find(self.layer_id)
+        layer = self._model.image_layers.updated.find(self.layer_id)
         assert layer is not None, "Control layer has been deleted"
         return layer
 

diff --git a/ai_diffusion/document.py b/ai_diffusion/document.py
@@ -316,6 +316,11 @@ def update(self):
     def find(self, id: QUuid):
         return next((l for l in self._layers if l.uniqueId() == id), None)
 
+    @property
+    def updated(self):
+        self.update()
+        return self
+
     def __iter__(self):
         return iter(self._layers)
 

diff --git a/ai_diffusion/ui/widget.py b/ai_diffusion/ui/widget.py
@@ -130,7 +130,6 @@ def __init__(self, model: Model, control: ControlLayer, parent: ControlListWidge
         self.add_pose_button = QToolButton(self)
         self.add_pose_button.setToolButtonStyle(Qt.ToolButtonStyle.ToolButtonIconOnly)
         self.add_pose_button.setIcon(theme.icon("add-pose"))
-        self.add_pose_button.setToolTip("Add new character pose to selected layer")
         self.add_pose_button.clicked.connect(self._add_pose_character)
 
         self.strength_spin = QSpinBox(self)
@@ -175,8 +174,10 @@ def __init__(self, model: Model, control: ControlLayer, parent: ControlListWidge
         control.is_supported_changed.connect(self._update_visibility)
         control.can_generate_changed.connect(self._update_visibility)
         control.show_end_changed.connect(self._update_visibility)
-        control.is_pose_vector_changed.connect(self._update_visibility)
+        control.mode_changed.connect(self._update_visibility)
+        control.is_pose_vector_changed.connect(self._update_pose_utils)
         self._update_visibility()
+        self._update_pose_utils()
 
     def _update_layers(self):
         layers: reversed[krita.Node] = reversed(self._model.image_layers)
@@ -202,7 +203,9 @@ def _update_visibility(self):
         def controls():
             self.layer_select.setVisible(self._control.is_supported)
             self.generate_button.setVisible(self._control.can_generate)
-            self.add_pose_button.setVisible(self._control.is_pose_vector)
+            self.add_pose_button.setVisible(
+                self._control.is_supported and self._control.mode is ControlMode.pose
+            )
             self.strength_spin.setVisible(self._control.is_supported)
             self.end_spin.setVisible(self._control.show_end)
 
@@ -216,6 +219,14 @@ def error():
             controls()
             error()
 
+    def _update_pose_utils(self):
+        self.add_pose_button.setEnabled(self._control.is_pose_vector)
+        self.add_pose_button.setToolTip(
+            "Add new character pose to selected layer"
+            if self._control.is_pose_vector
+            else "Disabled: selected layer must be a vector layer to add a pose"
+        )
+
     def _set_error(self, error: str):
         parts = error.split("[", 2)
         self.error_text.setText(parts[0])