Fix adding reference images in edit mode

Author: Acly

ai_diffusion/control.py

@@ -59,6 +59,7 @@ def __init__(self, model: model.Model, mode: ControlMode, layer_id: QUuid, index
 
         self.mode_changed.connect(self._update_is_supported)
         model.style_changed.connect(self._update_is_supported)
+        model.edit_mode_changed.connect(self._update_is_supported)
         root.connection.state_changed.connect(self._update_is_supported)
         self.layer_id_changed.connect(self._update_is_pose_vector)
         model.jobs.job_finished.connect(self._update_active_job)

ai_diffusion/model.py

@@ -184,7 +184,7 @@ def _generate(self, queue_mode: QueueMode):
         eventloop.run(_report_errors(self, jobs))
 
     def _prepare_workflow(self, dryrun=False):
-        is_edit = self.arch.is_edit or (self.edit_mode and self.edit_style is not None)
+        is_edit = self.arch.is_edit
         workflow_kind = WorkflowKind.generate
         if self.strength < 1.0 or is_edit:
             workflow_kind = WorkflowKind.refine
@@ -850,7 +850,7 @@ def prompt_translation_language(self):
 
     @property
     def arch(self):
-        return resolve_arch(self.style, self._connection.client_if_connected)
+        return resolve_arch(self.active_style, self._connection.client_if_connected)
 
     @property
     def history(self):
@@ -885,7 +885,8 @@ def name(self):
 
     @property
     def edit_style(self) -> Style | None:
-        if self.arch.is_edit:
+        style_arch = resolve_arch(self.style, self._connection.client_if_connected)
+        if style_arch.is_edit:
             return self.style
         if style_id := self.style.linked_edit_style:
             if style := Styles.list().find(style_id):

ai_diffusion/ui/generation.py

@@ -757,8 +757,8 @@ def model(self, model: Model):
                 model.region_only_changed.connect(self.update_generate_options),
                 model.style_changed.connect(self.update_generate_options),
                 model.edit_mode_changed.connect(self.update_generate_options),
-                self.add_control_button.clicked.connect(model.regions.add_control),
-                self.add_region_button.clicked.connect(model.regions.create_region_group),
+                self.add_control_button.clicked.connect(self.add_control),
+                self.add_region_button.clicked.connect(self.add_region),
                 self.region_prompt.activated.connect(model.generate),
                 self.generate_button.clicked.connect(model.generate),
                 self.generate_button.ctrl_clicked.connect(model.generate_replace),
@@ -846,7 +846,7 @@ def _create_edit_menu(self):
     def show_inpaint_menu(self):
         width = self.generate_button.width() + self.inpaint_mode_button.width()
         pos = QPoint(0, self.generate_button.height())
-        if self.model.arch.is_edit:
+        if not self.model.edit_mode and self.model.arch.is_edit:
             menu = self.edit_menu
         elif self.model.strength == 1.0:
             if self.model.region_only:
@@ -873,12 +873,18 @@ def toggle_region_only(self, checked: bool):
         self.model.region_only = checked
 
     def update_add_region_button(self):
-        supported = self.model.arch.supports_regions and not self.model.edit_mode
+        supported = self.model.arch.supports_regions
         self.add_region_button.setEnabled(supported)
         self.add_region_button.setToolTip(
             _("Add Region") if supported else _("Regions are not supported for the current model")
         )
 
+    def add_region(self):
+        self.model.active_regions.create_region_group()
+
+    def add_control(self):
+        self.model.active_regions.add_control()
+
     def update_generate_options(self):
         if not self.model.has_document:
             return
@@ -890,8 +896,10 @@ def update_generate_options(self):
         has_regions = len(regions) > 0
         has_active_region = regions.is_linked(self.model.layers.active)
         is_region_only = has_regions and has_active_region and self.model.region_only
-        is_edit = self.model.arch.is_edit or self.model.edit_mode
-        can_switch_edit = not self.model.arch.is_edit and self.model.edit_style is not None
+        is_edit = self.model.arch.is_edit
+        can_switch_edit = (
+            self.model.style.linked_edit_style != "" and self.model.edit_style is not None
+        )
         self.region_mask_button.setVisible(has_regions)
         self.region_mask_button.setEnabled(has_active_region)
         self.region_mask_button.setIcon(_region_mask_button_icons[is_region_only])