setiastrosuitepro 1.8.0.post3__py3-none-any.whl → 1.8.1.post2__py3-none-any.whl

setiastro/saspro/__main__.py

@@ -538,7 +538,18 @@ def _bootstrap_imports():
     add_runtime_to_sys_path(status_cb=lambda *_: None)
     _ban_shadow_torch_paths(status_cb=lambda *_: None)
     _purge_bad_torch_from_sysmodules(status_cb=lambda *_: None)
-
+    _update_splash(QCoreApplication.translate("Splash", "Preparing AI runtime cache..."), 7)
+    try:
+        from setiastro.saspro.runtime_torch import prewarm_torch_cache
+        prewarm_torch_cache(
+            status_cb=lambda *_: None,   # keep console clean during splash
+            require_torchaudio=True,
+            ensure_venv=True,
+            ensure_numpy=False,
+            validate_marker=True,
+        )
+    except Exception:
+        pass
     _update_splash(QCoreApplication.translate("Splash", "Loading standard libraries..."), 10)
 
     # ----------------------------------------

setiastro/saspro/_generated/build_info.py

@@ -1,3 +1,3 @@
 # Auto-generated at build time. Do not edit.
-BUILD_TIMESTAMP = "2026-01-26T19:14:34Z"
-APP_VERSION = "1.8.0.post3"
+BUILD_TIMESTAMP = "2026-01-28T14:34:03Z"
+APP_VERSION = "1.8.1.post2"

setiastro/saspro/cosmicclarity_engines/darkstar_engine.py

@@ -7,6 +7,7 @@ from typing import Any, Callable, Optional
 import numpy as np
 
 from setiastro.saspro.resources import get_resources
+from setiastro.saspro.runtime_torch import _user_runtime_dir, _venv_paths, _check_cuda_in_venv
 
 # Optional deps
 try:
@@ -44,13 +45,18 @@ def _autocast_context(torch, device) -> Any:
             major, minor = torch.cuda.get_device_capability()
             cap = float(f"{major}.{minor}")
             if cap >= 8.0:
-                # Preferred API (torch >= 1.10-ish; definitely in 2.x)
                 if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"):
                     return torch.amp.autocast(device_type="cuda")
-                # Fallback for older torch
                 return torch.cuda.amp.autocast()
+
+        elif hasattr(device, "type") and device.type == "mps":
+            # MPS often benefits from autocast in newer torch versions
+            if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"):
+                return torch.amp.autocast(device_type="mps")
+
     except Exception:
         pass
+
     return _nullcontext()
 
 
@@ -365,7 +371,21 @@ def load_darkstar_models(*, use_gpu: bool, color: bool, status_cb=print) -> Dark
         m = DarkStarModels(device=dev, is_onnx=False, model=net, torch=torch, chunk_size=512)
         _MODELS_CACHE[key] = m
         return m
+    # ---------------- MPS (torch) ----------------
+    if use_gpu and hasattr(torch, "backends") and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        backend_id = "mps"
+        key = (tag, backend_id)
+        if key in _MODELS_CACHE:
+            return _MODELS_CACHE[key]
 
+        dev = torch.device("mps")
+        status_cb("Dark Star: using MPS")
+        Net = _build_darkstar_torch_models(torch)
+        net = Net(pth, None).eval().to(dev)
+
+        m = DarkStarModels(device=dev, is_onnx=False, model=net, torch=torch, chunk_size=512)
+        _MODELS_CACHE[key] = m
+        return m
     # ---------------- DirectML (torch-directml) ----------------
     if use_gpu and is_windows:
         try:

setiastro/saspro/cosmicclarity_engines/denoise_engine.py

@@ -12,6 +12,7 @@ import cv2
 
 
 from setiastro.saspro.resources import get_resources
+from setiastro.saspro.runtime_torch import _user_runtime_dir, _venv_paths, _check_cuda_in_venv
 
 warnings.filterwarnings("ignore")
 
@@ -50,15 +51,19 @@ def _autocast_context(torch, device) -> Any:
             major, minor = torch.cuda.get_device_capability()
             cap = float(f"{major}.{minor}")
             if cap >= 8.0:
-                # Preferred API (torch >= 1.10-ish; definitely in 2.x)
                 if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"):
                     return torch.amp.autocast(device_type="cuda")
-                # Fallback for older torch
                 return torch.cuda.amp.autocast()
+
+        elif hasattr(device, "type") and device.type == "mps":
+            # MPS often benefits from autocast in newer torch versions
+            if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"):
+                return torch.amp.autocast(device_type="mps")
+
     except Exception:
         pass
-    return _nullcontext()
 
+    return _nullcontext()
 
 
 # ----------------------------
@@ -148,7 +153,18 @@ def load_models(use_gpu: bool = True, status_cb=print) -> Dict[str, Any]:
                 f"{'onnx' if models['is_onnx'] else 'torch'} / device={models['device']!r}")        
         _cached_models[key] = models
         return models
-
+    # >>> ADD THIS BLOCK HERE <<<
+    # 2) MPS (macOS Apple Silicon)
+    if use_gpu and hasattr(torch, "backends") and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = torch.device("mps")
+        status_cb("CosmicClarity Denoise: using MPS")
+        mono_model = _load_torch_model(torch, device, R.CC_DENOISE_PTH)
+        models = {"device": device, "is_onnx": False, "mono_model": mono_model, "torch": torch}
+        status_cb(f"Denoise backend resolved: "
+                f"{'onnx' if models['is_onnx'] else 'torch'} / device={models['device']!r}")
+        _cached_models[key] = models
+        return models
+    # >>> END INSERT <<<
     # 2) Torch-DirectML (Windows)
     if use_gpu and is_windows:
         try:
@@ -320,16 +336,19 @@ def denoise_chroma(cb: np.ndarray,
 def split_image_into_chunks_with_overlap(image, chunk_size, overlap):
     height, width = image.shape[:2]
     chunks = []
-    step_size = chunk_size - overlap  # Define how much to step over (overlapping area)
+    step_size = chunk_size - overlap
 
     for i in range(0, height, step_size):
         for j in range(0, width, step_size):
             end_i = min(i + chunk_size, height)
             end_j = min(j + chunk_size, width)
+            if end_i <= i or end_j <= j:
+                continue
             chunk = image[i:end_i, j:end_j]
-            chunks.append((chunk, i, j))  # Return chunk and its position
+            chunks.append((chunk, i, j))
     return chunks
 
+
 def blend_images(before, after, amount):
     return (1 - amount) * before + amount * after
 
@@ -337,29 +356,60 @@ def stitch_chunks_ignore_border(chunks, image_shape, border_size: int = 16):
     """
     chunks: list of (chunk, i, j) or (chunk, i, j, is_edge)
     image_shape: (H,W)
+    Robust to boundary clipping (prevents 256x256 -> 256x0 broadcasts).
     """
     H, W = image_shape
     stitched = np.zeros((H, W), dtype=np.float32)
     weights  = np.zeros((H, W), dtype=np.float32)
 
     for entry in chunks:
-        # accept both 3-tuple and 4-tuple
         if len(entry) == 3:
             chunk, i, j = entry
         else:
             chunk, i, j, _ = entry
 
         h, w = chunk.shape[:2]
+        if h <= 0 or w <= 0:
+            continue
+
         bh = min(border_size, h // 2)
         bw = min(border_size, w // 2)
 
+        # inner region in chunk coords
+        y0 = i + bh
+        y1 = i + h - bh
+        x0 = j + bw
+        x1 = j + w - bw
+
+        if y1 <= y0 or x1 <= x0:
+            continue
+
         inner = chunk[bh:h-bh, bw:w-bw]
-        stitched[i+bh:i+h-bh, j+bw:j+w-bw] += inner
-        weights[i+bh:i+h-bh, j+bw:j+w-bw] += 1.0
+
+        # clip destination to image bounds
+        yy0 = max(0, y0)
+        yy1 = min(H, y1)
+        xx0 = max(0, x0)
+        xx1 = min(W, x1)
+
+        if yy1 <= yy0 or xx1 <= xx0:
+            continue
+
+        # clip source to match clipped destination
+        sy0 = yy0 - y0
+        sy1 = sy0 + (yy1 - yy0)
+        sx0 = xx0 - x0
+        sx1 = sx0 + (xx1 - xx0)
+
+        src = inner[sy0:sy1, sx0:sx1]
+
+        stitched[yy0:yy1, xx0:xx1] += src
+        weights[yy0:yy1,  xx0:xx1] += 1.0
 
     stitched /= np.maximum(weights, 1.0)
     return stitched
 
+
 def replace_border(original_image, processed_image, border_size=16):
     # Ensure the dimensions of both images match
     if original_image.shape != processed_image.shape:
@@ -476,16 +526,19 @@ def denoise_channel(channel: np.ndarray, models: Dict[str, Any], *, progress_cb:
         original_chunk_shape = chunk.shape
 
         if is_onnx:
-            chunk_input = chunk[np.newaxis, np.newaxis, :, :].astype(np.float32)
-            chunk_input = np.tile(chunk_input, (1, 3, 1, 1))
-            if chunk_input.shape[2] != chunk_size or chunk_input.shape[3] != chunk_size:
+            h, w = original_chunk_shape  # <- the real chunk size
+
+            chunk_input = chunk[np.newaxis, np.newaxis, :, :].astype(np.float32)  # (1,1,h,w)
+            chunk_input = np.tile(chunk_input, (1, 3, 1, 1))                      # (1,3,h,w)
+
+            if h != chunk_size or w != chunk_size:
                 padded = np.zeros((1, 3, chunk_size, chunk_size), dtype=np.float32)
-                padded[:, :, :chunk_input.shape[2], :chunk_input.shape[3]] = chunk_input
+                padded[:, :, :h, :w] = chunk_input
                 chunk_input = padded
 
             input_name = model.get_inputs()[0].name
-            out = model.run(None, {input_name: chunk_input})[0]
-            denoised_chunk = out[0, 0, :original_chunk_shape[0], :original_chunk_shape[1]]
+            out = model.run(None, {input_name: chunk_input})[0]                   # (1,3,256,256) usually
+            denoised_chunk = out[0, 0, :h, :w]
 
         else:
             torch = models["torch"]

setiastro/saspro/cosmicclarity_engines/satellite_engine.py

@@ -46,15 +46,19 @@ def _autocast_context(torch, device) -> Any:
             major, minor = torch.cuda.get_device_capability()
             cap = float(f"{major}.{minor}")
             if cap >= 8.0:
-                # Preferred API (torch >= 1.10-ish; definitely in 2.x)
                 if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"):
                     return torch.amp.autocast(device_type="cuda")
-                # Fallback for older torch
                 return torch.cuda.amp.autocast()
+
+        elif hasattr(device, "type") and device.type == "mps":
+            # MPS often benefits from autocast in newer torch versions
+            if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"):
+                return torch.amp.autocast(device_type="mps")
+
     except Exception:
         pass
-    return _nullcontext()
 
+    return _nullcontext()
 
 
 # ----------------------------