setiastrosuitepro 1.6.4__py3-none-any.whl → 1.7.1.post2__py3-none-any.whl

setiastro/saspro/aberration_ai.py

@@ -7,13 +7,37 @@ import numpy as np
 import sys
 import platform  # add
 import time
+import subprocess
 
 IS_APPLE_ARM = (sys.platform == "darwin" and platform.machine() == "arm64")
 
+def _has_nvidia_gpu() -> bool:
+    """Check if system has an NVIDIA GPU (Linux/Windows)."""
+    try:
+        if platform.system() == "Linux":
+            r = subprocess.run(["nvidia-smi", "-L"], capture_output=True, timeout=2)
+            return "GPU" in (r.stdout.decode("utf-8", errors="ignore") or "")
+        elif platform.system() == "Windows":
+            try:
+                ps = subprocess.run(
+                    ["powershell", "-NoProfile", "-Command",
+                     "(Get-CimInstance Win32_VideoController | Select-Object -ExpandProperty Name) -join ';'"],
+                    capture_output=True, timeout=2
+                )
+                out = (ps.stdout.decode("utf-8", errors="ignore") or "").lower()
+                return "nvidia" in out
+            except Exception:
+                w = subprocess.run(["wmic", "path", "win32_VideoController", "get", "name"],
+                                  capture_output=True, timeout=2)
+                return "nvidia" in (w.stdout.decode("utf-8", errors="ignore") or "").lower()
+    except Exception:
+        pass
+    return False
+
 from PyQt6.QtCore import Qt, QThread, pyqtSignal, QStandardPaths, QSettings
 from PyQt6.QtWidgets import (
     QDialog, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QFileDialog,
-    QComboBox, QSpinBox, QProgressBar, QMessageBox, QCheckBox
+    QComboBox, QSpinBox, QProgressBar, QMessageBox, QCheckBox, QLineEdit, QApplication
 )
 from PyQt6.QtGui import QIcon
 from setiastro.saspro.config import Config
@@ -145,10 +169,13 @@ def _restore_output(arr: np.ndarray, channels_last: bool, was_uint16: bool, H: i
         arr = arr[0]                 # (H,W)
     return arr
 
-def run_onnx_tiled(session, img: np.ndarray, patch_size=512, overlap=64, progress_cb=None) -> np.ndarray:
+def run_onnx_tiled(session, img: np.ndarray, patch_size=512, overlap=64,
+                  progress_cb=None, cancel_cb=None) -> np.ndarray:
     """
     session: onnxruntime.InferenceSession
     img: mono (H,W) or RGB (H,W,3) numpy array
+
+    cancel_cb: callable -> bool, return True to cancel
     """
     arr, channels_last, was_uint16 = _prepare_input(img)      # (C,H,W)
     arr, H0, W0 = _pad_C_HW(arr, patch_size)
@@ -168,11 +195,15 @@ def run_onnx_tiled(session, img: np.ndarray, patch_size=512, overlap=64, progres
     for c in range(C):
         for i in hs:
             for j in ws:
-                patch = arr[c:c+1, i:i+patch_size, j:j+patch_size]  # (1, P, P)
+                if cancel_cb and cancel_cb():
+                    raise RuntimeError("Canceled")
+
+                patch = arr[c:c+1, i:i+patch_size, j:j+patch_size]  # (1,P,P)
                 inp = np.ascontiguousarray(patch[np.newaxis, ...], dtype=np.float32)  # (1,1,P,P)
 
                 out_patch = session.run(None, {inp_name: inp})[0]   # (1,1,P,P)
                 out_patch = np.squeeze(out_patch, axis=0)           # (1,P,P)
+
                 out[c:c+1, i:i+patch_size, j:j+patch_size] += out_patch * win
                 wgt[c:c+1, i:i+patch_size, j:j+patch_size] += win
 
@@ -184,7 +215,6 @@ def run_onnx_tiled(session, img: np.ndarray, patch_size=512, overlap=64, progres
     arr = out / wgt
     return _restore_output(arr, channels_last, was_uint16, H0, W0)
 
-
 # ---------- providers ----------
 def pick_providers(auto_gpu=True) -> list[str]:
     """
@@ -248,9 +278,11 @@ def _preserve_border(dst: np.ndarray, src: np.ndarray, px: int = 10) -> np.ndarr
 
 # ---------- worker ----------
 class _ONNXWorker(QThread):
-    progressed = pyqtSignal(int)         # 0..100
-    failed     = pyqtSignal(str)
-    finished_ok= pyqtSignal(np.ndarray)
+    progressed   = pyqtSignal(int)          # 0..100
+    failed       = pyqtSignal(str)
+    finished_ok  = pyqtSignal(np.ndarray)
+    canceled     = pyqtSignal()
+    log_message  = pyqtSignal(str)          # for console logging
 
     def __init__(self, model_path: str, image: np.ndarray, patch: int, overlap: int, providers: list[str]):
         super().__init__()
@@ -260,33 +292,115 @@ class _ONNXWorker(QThread):
         self.overlap = overlap
         self.providers = providers
         self.used_provider = None
+        self._cancel = False  # cooperative flag
+
+    def cancel(self):
+        # Safe to call from UI thread
+        self._cancel = True
+        self.requestInterruption()
+
+    def _is_canceled(self) -> bool:
+        return self._cancel or self.isInterruptionRequested()
 
     def run(self):
         if ort is None:
             self.failed.emit("onnxruntime is not installed.")
             return
+
+        # If canceled before start, exit cleanly
+        if self._is_canceled():
+            self.canceled.emit()
+            return
+
+        # Log available providers for debugging
+        avail_providers = ort.get_available_providers()
+        gpu_providers = [p for p in self.providers if p != "CPUExecutionProvider"]
+        has_nvidia = _has_nvidia_gpu()
+        
+        self.log_message.emit(f"🔍 Available ONNX providers: {', '.join(avail_providers)}")
+        self.log_message.emit(f"🔍 Attempting providers: {', '.join(self.providers)}")
+        print(f"🔍 Available ONNX providers: {', '.join(avail_providers)}")
+        print(f"🔍 Attempting providers: {', '.join(self.providers)}")
+        
+        # Check if NVIDIA GPU is present but CUDA provider is missing
+        if has_nvidia and "CUDAExecutionProvider" not in avail_providers:
+            msg = ("⚠️ GPU NVIDIA détecté mais CUDAExecutionProvider n'est pas disponible.\n"
+                   "   Vous devez installer 'onnxruntime-gpu' au lieu de 'onnxruntime'.\n"
+                   "   Commande: pip uninstall onnxruntime && pip install onnxruntime-gpu")
+            self.log_message.emit(msg)
+            print(msg)
+        
         try:
             sess = ort.InferenceSession(self.model_path, providers=self.providers)
             self.used_provider = (sess.get_providers()[0] if sess.get_providers() else None)
-        except Exception:
+            # Log successful GPU usage
+            if self.used_provider != "CPUExecutionProvider" and gpu_providers:
+                msg = f"✅ Aberration AI: Using GPU provider {self.used_provider}"
+                self.log_message.emit(msg)
+                print(msg)
+            elif has_nvidia and self.used_provider == "CPUExecutionProvider":
+                msg = ("⚠️ GPU NVIDIA détecté mais utilisation du CPU.\n"
+                       "   Installez 'onnxruntime-gpu' pour utiliser le GPU.")
+                self.log_message.emit(msg)
+                print(msg)
+        except Exception as e:
+            # Log the actual error for debugging
+            error_msg = str(e)
+            msg = f"⚠️ Aberration AI: GPU provider failed: {error_msg}"
+            self.log_message.emit(msg)
+            print(msg)
+            self.log_message.emit(f"Available providers: {', '.join(avail_providers)}")
+            print(f"Available providers: {', '.join(avail_providers)}")
+            self.log_message.emit(f"Attempted providers: {', '.join(self.providers)}")
+            print(f"Attempted providers: {', '.join(self.providers)}")
+            
+            # Check if onnxruntime-gpu is installed (CUDA provider should be available if it is)
+            if "CUDAExecutionProvider" in self.providers and "CUDAExecutionProvider" not in avail_providers:
+                if has_nvidia:
+                    msg = ("❌ CUDAExecutionProvider non disponible alors qu'un GPU NVIDIA est présent.\n"
+                           "   Installez 'onnxruntime-gpu': pip uninstall onnxruntime && pip install onnxruntime-gpu")
+                else:
+                    msg = "⚠️ CUDAExecutionProvider not available. You may need to install onnxruntime-gpu instead of onnxruntime."
+                self.log_message.emit(msg)
+                print(msg)
+            
             # fallback CPU if GPU fails
             try:
                 sess = ort.InferenceSession(self.model_path, providers=["CPUExecutionProvider"])
-                self.used_provider = "CPUExecutionProvider"  # NEW
+                self.used_provider = "CPUExecutionProvider"
+                msg = f"⚠️ Aberration AI: Falling back to CPU (GPU initialization failed: {error_msg})"
+                self.log_message.emit(msg)
+                print(msg)
             except Exception as e2:
-                self.failed.emit(f"Failed to init ONNX session:\n{e2}")
+                self.failed.emit(f"Failed to init ONNX session:\nGPU error: {error_msg}\nCPU error: {e2}")
                 return
 
         def cb(frac):
             self.progressed.emit(int(frac * 100))
 
         try:
-            out = run_onnx_tiled(sess, self.image, self.patch, self.overlap, cb)
+            out = run_onnx_tiled(
+                sess,
+                self.image,
+                self.patch,
+                self.overlap,
+                progress_cb=cb,
+                cancel_cb=self._is_canceled,
+            )
         except Exception as e:
-            self.failed.emit(str(e)); return
+            # Normalize cancel
+            msg = str(e) or "Error"
+            if "Canceled" in msg:
+                self.canceled.emit()
+            else:
+                self.failed.emit(msg)
+            return
 
-        self.finished_ok.emit(out)
+        if self._is_canceled():
+            self.canceled.emit()
+            return
 
+        self.finished_ok.emit(out)
 
 # ---------- dialog ----------
 class AberrationAIDialog(QDialog):
@@ -315,11 +429,29 @@ class AberrationAIDialog(QDialog):
         row.addWidget(QLabel(self.tr("Model:")))
         self.model_label = QLabel("—")
         self.model_label.setToolTip("")
-        btn_browse = QPushButton(self.tr("Browse…")); btn_browse.clicked.connect(self._browse_model)
+        btn_browse = QPushButton(self.tr("Browse…")); btn_browse.clicked.connect(self._browse_active_model)
         row.addWidget(self.model_label, 1)
         row.addWidget(btn_browse)
         v.addLayout(row)
+        # Custom model row (NEW)
+        row_custom = QHBoxLayout()
+        self.chk_use_custom = QCheckBox(self.tr("Use custom model file"))
+        self.chk_use_custom.setChecked(False)
+        self.chk_use_custom.toggled.connect(self._on_use_custom_toggled)
+
+        self.le_custom_model = QLineEdit()
+        self.le_custom_model.setReadOnly(True)
+        self.le_custom_model.setPlaceholderText(self.tr("No custom model selected"))
+        self.le_custom_model.setToolTip("")
+
+        btn_custom_clear = QPushButton(self.tr("Clear"))
+        btn_custom_clear.clicked.connect(self._clear_custom_model)
+
+        row_custom.addWidget(self.chk_use_custom)
+        row_custom.addWidget(self.le_custom_model, 1)
 
+        row_custom.addWidget(btn_custom_clear)
+        v.addLayout(row_custom)
         # Providers row
         row2 = QHBoxLayout()
         self.chk_auto = QCheckBox(self.tr("Auto GPU (if available)"))
@@ -373,7 +505,9 @@ class AberrationAIDialog(QDialog):
         self._model_path = None
         self._refresh_providers()
         self._load_last_model_from_settings()
-
+        self._load_last_custom_model_from_settings()
+        use_custom = QSettings().value("AberrationAI/use_custom_model", False, type=bool)
+        self.chk_use_custom.setChecked(bool(use_custom))
         if IS_APPLE_ARM:
             self.chk_auto.setChecked(False)
             self.chk_auto.setEnabled(False)
@@ -395,11 +529,73 @@ class AberrationAIDialog(QDialog):
         if p and os.path.isfile(p):
             self._set_model_path(p)
 
-    def _browse_model(self):
-        start_dir = _app_model_dir()
+    def _browse_active_model(self):
+        """
+        Single Browse button.
+        - If user picks a file inside the app model folder -> treat as "downloaded" selection (use_custom_model=False)
+        - If user picks a file outside -> treat as "custom" (use_custom_model=True)
+        """
+        app_dir = os.path.abspath(_app_model_dir())
+
+        # Start in last-used folder if possible
+        last_custom = QSettings().value("AberrationAI/custom_model_path", type=str) or ""
+        last_downloaded = QSettings().value("AberrationAI/model_path", type=str) or ""
+        start_dir = None
+        for candidate in (last_custom, last_downloaded):
+            if candidate and os.path.isfile(candidate):
+                d = os.path.dirname(candidate)
+                if os.path.isdir(d):
+                    start_dir = d
+                    break
+        if start_dir is None:
+            start_dir = app_dir
+
         p, _ = QFileDialog.getOpenFileName(self, "Select ONNX model", start_dir, "ONNX (*.onnx)")
-        if p:
-            self._set_model_path(p)
+        if not p:
+            return
+
+        p_abs = os.path.abspath(p)
+        # Determine if picked file is inside app model folder
+        in_app_dir = False
+        try:
+            in_app_dir = os.path.commonpath([app_dir, p_abs]) == app_dir
+        except Exception:
+            in_app_dir = p_abs.startswith(app_dir)
+
+        if in_app_dir:
+            # "Downloaded" selection
+            self._set_model_path(p_abs)
+            self._set_custom_model_path(None)
+            QSettings().setValue("AberrationAI/use_custom_model", False)
+            if hasattr(self, "chk_use_custom"):
+                self.chk_use_custom.setChecked(False)
+        else:
+            # "Custom" selection
+            self._set_custom_model_path(p_abs)
+            QSettings().setValue("AberrationAI/use_custom_model", True)
+            if hasattr(self, "chk_use_custom"):
+                self.chk_use_custom.setChecked(True)
+
+        # Keep visuals in sync
+        self._refresh_model_label()
+        self._refresh_custom_row_visibility()
+
+
+    def _refresh_model_label(self):
+        downloaded = QSettings().value("AberrationAI/model_path", type=str) or ""
+        custom     = QSettings().value("AberrationAI/custom_model_path", type=str) or ""
+        use_custom = QSettings().value("AberrationAI/use_custom_model", False, type=bool)
+
+        if use_custom and custom:
+            self.model_label.setText(f"Custom: {os.path.basename(custom)}")
+            self.model_label.setToolTip(custom)
+        elif downloaded:
+            self.model_label.setText(f"Downloaded: {os.path.basename(downloaded)}")
+            self.model_label.setToolTip(downloaded)
+        else:
+            self.model_label.setText("—")
+            self.model_label.setToolTip("")
+
 
     def _open_model_folder(self):
         d = _app_model_dir()
@@ -412,6 +608,108 @@ class AberrationAIDialog(QDialog):
                 import subprocess; subprocess.Popen(["xdg-open", d])
         except Exception:
             webbrowser.open(f"file://{d}")
+    # ----- custom model helpers (NEW) -----
+    def _set_custom_model_path(self, p: str | None):
+        if p:
+            self.le_custom_model.setText(os.path.basename(p))
+            self.le_custom_model.setToolTip(p)
+            QSettings().setValue("AberrationAI/custom_model_path", p)
+        else:
+            self.le_custom_model.clear()
+            self.le_custom_model.setToolTip("")
+            QSettings().remove("AberrationAI/custom_model_path")
+
+    def _load_last_custom_model_from_settings(self):
+        p = QSettings().value("AberrationAI/custom_model_path", type=str)
+        if p:
+            if os.path.isfile(p):
+                self._set_custom_model_path(p)
+            else:
+                # Keep the broken path visible in tooltip for debugging
+                if hasattr(self, "le_custom_model"):
+                    self.le_custom_model.setText(os.path.basename(p) + "  (missing)")
+                    self.le_custom_model.setToolTip(p)
+
+        # After both loads, sync labels/visibility
+        self._refresh_model_label()
+        self._refresh_custom_row_visibility()
+
+    def _refresh_custom_row_visibility(self):
+        """
+        If you keep the custom row in the UI, hide the path field unless custom is enabled.
+        """
+        if not hasattr(self, "le_custom_model"):
+            return
+        use_custom = QSettings().value("AberrationAI/use_custom_model", False, type=bool)
+        self.le_custom_model.setVisible(bool(use_custom))
+
+
+    def _refresh_model_label(self):
+        downloaded = QSettings().value("AberrationAI/model_path", type=str) or ""
+        custom     = QSettings().value("AberrationAI/custom_model_path", type=str) or ""
+        use_custom = QSettings().value("AberrationAI/use_custom_model", False, type=bool)
+
+        # Prefer custom only if enabled AND the file exists
+        if use_custom and custom:
+            if os.path.isfile(custom):
+                self.model_label.setText(f"Custom: {os.path.basename(custom)}")
+                self.model_label.setToolTip(custom)
+                return
+            else:
+                self.model_label.setText(f"Custom: {os.path.basename(custom)}  (missing)")
+                self.model_label.setToolTip(custom)
+                return
+
+        # Otherwise show downloaded if valid
+        if downloaded and os.path.isfile(downloaded):
+            self.model_label.setText(f"Downloaded: {os.path.basename(downloaded)}")
+            self.model_label.setToolTip(downloaded)
+        else:
+            self.model_label.setText("—")
+            self.model_label.setToolTip("")
+
+
+    def _browse_custom_model(self):
+        # Start at last dir if possible, else app model dir
+        last = QSettings().value("AberrationAI/custom_model_path", type=str) or ""
+        start_dir = os.path.dirname(last) if last and os.path.isdir(os.path.dirname(last)) else _app_model_dir()
+        p, _ = QFileDialog.getOpenFileName(self, "Select custom ONNX model", start_dir, "ONNX (*.onnx)")
+        if p:
+            self._set_custom_model_path(p)
+            QSettings().setValue("AberrationAI/use_custom_model", True)
+            if not self.chk_use_custom.isChecked():
+                self.chk_use_custom.setChecked(True)
+
+    def _clear_custom_model(self):
+        self._set_custom_model_path(None)
+        QSettings().setValue("AberrationAI/use_custom_model", False)
+        if hasattr(self, "chk_use_custom"):
+            self.chk_use_custom.setChecked(False)
+
+        self._refresh_model_label()
+        self._refresh_custom_row_visibility()
+
+
+    def _on_use_custom_toggled(self, on: bool):
+        QSettings().setValue("AberrationAI/use_custom_model", bool(on))
+
+        if on:
+            p = QSettings().value("AberrationAI/custom_model_path", type=str) or ""
+            if not (p and os.path.isfile(p)):
+                # Don’t spawn another browse button path; use the ONE browse if they want
+                QMessageBox.information(
+                    self,
+                    self.tr("Custom model"),
+                    self.tr("Custom model is enabled, but no custom file is selected.\n"
+                            "Click Browse… to choose a model file.")
+                )
+                # Optional: auto-open the single browse:
+                # self._browse_active_model()
+                # return
+
+        self._refresh_model_label()
+        self._refresh_custom_row_visibility()
+
 
     # ----- provider UI -----
     def _log(self, msg: str):  # NEW
@@ -477,8 +775,16 @@ class AberrationAIDialog(QDialog):
     def _on_download_ok(self, path: str):
         self.progress.setValue(100)
         self._set_model_path(path)
+
+        # Download becomes the active model unless custom is explicitly enabled
+        if not QSettings().value("AberrationAI/use_custom_model", False, type=bool):
+            self._set_custom_model_path(None)
+
         QMessageBox.information(self, "Model", f"Downloaded: {os.path.basename(path)}")
 
+        self._refresh_model_label()
+        self._refresh_custom_row_visibility()
+
     # ----- run -----
     def _run(self):
         if ort is None:
@@ -489,7 +795,22 @@ class AberrationAIDialog(QDialog):
                 "Please try installing an earlier version (for example 1.19.x) and try again."
             )
             return
-        if not self._model_path or not os.path.isfile(self._model_path):
+        
+        # Choose model path (normal vs custom)
+        use_custom = QSettings().value("AberrationAI/use_custom_model", False, type=bool)
+        downloaded = QSettings().value("AberrationAI/model_path", type=str) or ""
+        custom     = QSettings().value("AberrationAI/custom_model_path", type=str) or ""
+
+        model_path = custom if use_custom else downloaded
+        if self.chk_use_custom.isChecked():
+            cp = QSettings().value("AberrationAI/custom_model_path", type=str)
+            if cp and os.path.isfile(cp):
+                model_path = cp
+            else:
+                QMessageBox.warning(self, "Model", "Custom model is enabled but the file is missing. Please browse to a valid .onnx.")
+                return
+
+        if not model_path or not os.path.isfile(model_path):
             QMessageBox.warning(self, "Model", "Please select or download a valid .onnx model first.")
             return
 
@@ -516,7 +837,7 @@ class AberrationAIDialog(QDialog):
                 providers = [sel] if sel else ["CPUExecutionProvider"]
 
         # --- make patch match the model's requirement (if fixed) ---
-        req = _model_required_patch(self._model_path)
+        req = _model_required_patch(model_path)
         if req and req > 0:
             patch = req
             try:
@@ -537,26 +858,34 @@ class AberrationAIDialog(QDialog):
 
         self._t_start = time.perf_counter()
         prov_txt = ("auto" if self.chk_auto.isChecked() else self.cmb_provider.currentText() or "CPU")
-        self._log(f"🚀 Aberration AI: model={os.path.basename(self._model_path)}, "
+        self._log(f"🚀 Aberration AI: model={os.path.basename(model_path)}, "
                   f"provider={prov_txt}, patch={patch}, overlap={overlap}")
+        
+        self._effective_model_path = model_path
 
         # -------- run worker --------
         self.progress.setValue(0)
         self.btn_run.setEnabled(False)
 
-        self._worker = _ONNXWorker(self._model_path, img, patch, overlap, providers)
+        self._worker = _ONNXWorker(model_path, img, patch, overlap, providers)
         self._worker.progressed.connect(self.progress.setValue)
         self._worker.failed.connect(self._on_failed)
         self._worker.finished_ok.connect(self._on_ok)
         self._worker.finished.connect(self._on_worker_finished)
+        self._worker.log_message.connect(self._log)  # Connect log messages to console
         self._worker.start()
 
 
     def _on_failed(self, msg: str):
-        self._log(f"❌ Aberration AI failed: {msg}")   # NEW
+        model_path = getattr(self, "_effective_model_path", self._model_path)
+        self._log(f"❌ Aberration AI failed: {msg}")
         QMessageBox.critical(self, "ONNX Error", msg)
+        self.reject()   # closes the dialog
 
     def _on_ok(self, out: np.ndarray):
+        used = getattr(self._worker, "used_provider", None) or \
+            (self.cmb_provider.currentText() if not self.chk_auto.isChecked() else "auto")        
+        model_path = getattr(self, "_effective_model_path", self._model_path)
         doc = self.get_active_doc()
         if doc is None or getattr(doc, "image", None) is None:
             QMessageBox.warning(self, "Image", "No active image.")
@@ -578,11 +907,10 @@ class AberrationAIDialog(QDialog):
             "processing_parameters": {
                 **(getattr(doc, "metadata", {}) or {}).get("processing_parameters", {}),
                 "AberrationAI": {
-                    "model_path": self._model_path,
+                    "model_path": model_path,
                     "patch_size": int(self.spin_patch.value()),
                     "overlap": int(self.spin_overlap.value()),
-                    "provider": (self.cmb_provider.currentText()
-                                if not self.chk_auto.isChecked() else "auto"),
+                    "provider": used,
                     "border_px": BORDER_PX,
                 }
             }
@@ -615,7 +943,7 @@ class AberrationAIDialog(QDialog):
             if main is not None:
                 auto_gpu = bool(self.chk_auto.isChecked())
                 preset = {
-                    "model": self._model_path,
+                    "model": model_path,
                     "patch": int(self.spin_patch.value()),
                     "overlap": int(self.spin_overlap.value()),
                     "border_px": int(BORDER_PX),
@@ -674,21 +1002,24 @@ class AberrationAIDialog(QDialog):
         BORDER_PX = 10  # same value used above
         self._log(
             f"✅ Aberration AI applied "
-            f"(model={os.path.basename(self._model_path)}, provider={used}, "
+            f"(model={os.path.basename(model_path)}, provider={used}, "
             f"patch={int(self.spin_patch.value())}, overlap={int(self.spin_overlap.value())}, "
             f"border={BORDER_PX}px, time={dt:.2f}s)"
         )
 
         self.progress.setValue(100)
-        # Dialog stays open so user can apply to other images
+        # NEW: close this UI after a successful run
+        self.accept()   # or self.close()
+        return
 
     def _on_worker_finished(self):
-        # If dialog is already gone, this method is never called because the receiver (self)
-        # has been destroyed and Qt auto-disconnects the signal.
+        # Dialog might have been closed by _on_ok()
+        if not self.isVisible():
+            return
+
         if hasattr(self, "btn_run"):
             try:
                 self.btn_run.setEnabled(True)
             except RuntimeError:
-                # Button already deleted; ignore
                 pass
         self._worker = None

setiastro/saspro/aberration_ai_preset.py

@@ -4,7 +4,7 @@ import os
 import time
 import numpy as np
 from PyQt6.QtCore import QTimer
-from PyQt6.QtWidgets import QDialog, QVBoxLayout, QProgressBar, QPushButton, QMessageBox, QFormLayout, QDialogButtonBox, QSpinBox, QCheckBox, QComboBox, QLabel
+from PyQt6.QtWidgets import QDialog, QVBoxLayout, QProgressBar, QPushButton, QMessageBox, QFormLayout, QDialogButtonBox, QSpinBox, QCheckBox, QComboBox, QLabel, QApplication
 
 from PyQt6.QtCore import QSettings
 # reuse everything from the UI module
@@ -91,13 +91,29 @@ def run_aberration_ai_via_preset(main, preset: dict | None = None, doc=None):
     worker = _ONNXWorker(model, img, patch, overlap, providers)
     worker.progressed.connect(bar.setValue)
 
+    def _cancel_clicked():
+        btn.setEnabled(False)
+        btn.setText("Canceling…")
+        worker.cancel()  # <-- SAFE
+        QApplication.processEvents()
+
     def _fail(msg: str):
         try:
             if hasattr(main, "_log"):
                 main._log(f"❌ Aberration AI failed: {msg}")
         except Exception:
             pass
-        QMessageBox.critical(main, "Aberration AI", msg)
+        # If canceled, don't pop an error box
+        if "Canceled" not in (msg or ""):
+            QMessageBox.critical(main, "Aberration AI", msg)
+        dlg.close()
+
+    def _canceled():
+        try:
+            if hasattr(main, "_log"):
+                main._log("⛔ Aberration AI canceled.")
+        except Exception:
+            pass
         dlg.close()
 
     def _ok(out: np.ndarray):
@@ -157,13 +173,23 @@ def run_aberration_ai_via_preset(main, preset: dict | None = None, doc=None):
             dlg.close()
 
     worker.failed.connect(_fail)
+    worker.canceled.connect(_canceled)          # <-- NEW
     worker.finished_ok.connect(_ok)
     worker.finished.connect(lambda: btn.setEnabled(False))
-    btn.clicked.connect(worker.terminate)
+
+    btn.clicked.connect(_cancel_clicked)
+
+    # If user closes dialog via window X, also cancel
+    dlg.rejected.connect(_cancel_clicked)
 
     worker.start()
     dlg.exec()
 
+    # Ensure the worker is not left running after the modal closes
+    if worker.isRunning():
+        worker.cancel()
+        worker.wait(2000)  # don't hang forever; just give it a moment
+
     # clear the guard after a brief tick so downstream signals don’t re-open UI
     def _clear():
         for k in ("_aberration_ai_headless_running", "_aberration_ai_guard"):